





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           =			logging.get_logger(__name__)

_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""",
    """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           =			[
    """ctc_proj""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]


def _A       (   __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__	):
   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
            lowercase__			       =  "lm_head"

      lowercase__			       =  getattr(__magic_name__  , __magic_name__	)

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

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


def _A       (   __magic_name__  , __magic_name__  , __magic_name__	):
   lowercase__			       =  []
   lowercase__			       =  fairseq_model.state_dict()

   lowercase__			       =  hf_model.unispeech.feature_extractor

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

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


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


@torch.no_grad()
def _A       (   __magic_name__  , __magic_name__  , __magic_name__=None  , __magic_name__=None  , __magic_name__=True	):
   if config_path is not None:
      lowercase__			       =  UniSpeechConfig.from_pretrained(__magic_name__	)
   else:
      lowercase__			       =  UniSpeechConfig()

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

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

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

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

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

   lowercase__			       =  model[0].eval()

   recursively_load_weights(__magic_name__  , __magic_name__  , __magic_name__	)

   hf_unispeech.save_pretrained(__magic_name__	)


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













from __future__ import annotations


def _A       (   __magic_name__  , __magic_name__	):
   lowercase__			       =  []
   create_all_state(1  , __magic_name__  , __magic_name__  , []  , __magic_name__	)
   return result


def _A       (   __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , ):
   if level == 0:
      total_list.append(current_list[:]	)
      return

   for i in range(__magic_name__  , total_number - level + 2	):
      current_list.append(__magic_name__	)
      create_all_state(i + 1  , __magic_name__  , level - 1  , __magic_name__  , __magic_name__	)
      current_list.pop()


def _A       (   __magic_name__	):
   for i in total_list:
      print(*__magic_name__	)


if __name__ == "__main__":
  _snake_case           =			4
  _snake_case           =			2
  _snake_case           =			generate_all_combinations(n, k)
  print_all_state(total_list)






