





import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, 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',
    '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',
}
snake_case_           =      [
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'label_embeddings_concat',
    'speaker_proj',
    'layer_norm_for_extract',
]






def lowerCamelCase__   (							snake_case_     :  int	,				snake_case_     :  str	,				snake_case_     :  Optional[int]	,				snake_case_     :  Any	,				snake_case_     :  List[str]	)				->    List[str]:
   for attribute in key.split('''.'''	):
      __snake_case								=       getattr(snake_case_	,				snake_case_	)

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

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






def lowerCamelCase__   (							snake_case_     :  List[str]	,				snake_case_     :  List[str]	)				->    Optional[int]:
   __snake_case								=       []
   __snake_case								=       fairseq_model.state_dict()

   __snake_case								=       hf_model.unispeech_sat.feature_extractor

   for name, value in fairseq_dict.items():
      __snake_case								=       False
      if "conv_layers" in name:
         load_conv_layer(
             snake_case_	,				snake_case_	,				snake_case_	,				snake_case_	,				hf_model.config.feat_extract_norm == '''group'''	,				)
         __snake_case								=       True
      else:
         for key, mapped_key in MAPPING.items():
            __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
               __snake_case								=       True
               if "*" in mapped_key:
                  __snake_case								=       name.split(snake_case_	)[0].split('''.'''	)[-2]
                  __snake_case								=       mapped_key.replace('''*'''	,				snake_case_	)
               if "weight_g" in name:
                  __snake_case								=       '''weight_g'''
               elif "weight_v" in name:
                  __snake_case								=       '''weight_v'''
               elif "bias" in name:
                  __snake_case								=       '''bias'''
               elif "weight" in name:
                  # TODO: don't match quantizer.weight_proj
                  __snake_case								=       '''weight'''
               else:
                  __snake_case								=       None
               set_recursively(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 lowerCamelCase__   (							snake_case_     :  Union[str, Any]	,				snake_case_     :  List[Any]	,				snake_case_     :  Tuple	,				snake_case_     :  Dict	,				snake_case_     :  Optional[Any]	)				->    int:
   __snake_case								=       full_name.split('''conv_layers.'''	)[-1]
   __snake_case								=       name.split('''.'''	)
   __snake_case								=       int(items[0]	)
   __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."""	)
         __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."""	)
         __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."""	)
         __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."""	)
         __snake_case								=       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 lowerCamelCase__   (							snake_case_     :  Dict	,				snake_case_     :  List[str]	,				snake_case_     :  Union[str, Any]=None	,				snake_case_     :  List[Any]=None	,				snake_case_     :  int=True	)				->    Dict:
   if config_path is not None:
      __snake_case								=       UniSpeechSatConfig.from_pretrained(snake_case_	)
   else:
      __snake_case								=       UniSpeechSatConfig()

   __snake_case								=       ''''''

   if is_finetuned:
      __snake_case								=       UniSpeechSatForCTC(snake_case_	)
   else:
      __snake_case								=       UniSpeechSatForPreTraining(snake_case_	)

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

   recursively_load_weights(snake_case_	,				snake_case_	)

   hf_wavavec.save_pretrained(snake_case_	)


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_sat_checkpoint(
        args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
    )








import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
snake_case_           =      data_utils.TransfoXLTokenizer
snake_case_           =      data_utils.TransfoXLCorpus
snake_case_           =      data_utils
snake_case_           =      data_utils






def lowerCamelCase__   (							snake_case_     :  Union[str, Any]	,				snake_case_     :  int	,				snake_case_     :  List[Any]	,				snake_case_     :  int	)				->    Dict:
   if transfo_xl_dataset_file:
      # Convert a pre-processed corpus (see original TensorFlow repo)
      with open(snake_case_	,				'''rb'''	) as fp:
         __snake_case								=       pickle.load(snake_case_	,				encoding='''latin1'''	)
      # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
      __snake_case								=       pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
      print(f"""Save vocabulary to {pytorch_vocab_dump_path}"""	)
      __snake_case								=       corpus.vocab.__dict__
      torch.save(snake_case_	,				snake_case_	)

      __snake_case								=       corpus.__dict__
      corpus_dict_no_vocab.pop('''vocab'''	,				snake_case_	)
      __snake_case								=       pytorch_dump_folder_path + '''/''' + CORPUS_NAME
      print(f"""Save dataset to {pytorch_dataset_dump_path}"""	)
      torch.save(snake_case_	,				snake_case_	)

   if tf_checkpoint_path:
      # Convert a pre-trained TensorFlow model
      __snake_case								=       os.path.abspath(snake_case_	)
      __snake_case								=       os.path.abspath(snake_case_	)

      print(f"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""	)
      # Initialise PyTorch model
      if transfo_xl_config_file == "":
         __snake_case								=       TransfoXLConfig()
      else:
         __snake_case								=       TransfoXLConfig.from_json_file(snake_case_	)
      print(f"""Building PyTorch model from configuration: {config}"""	)
      __snake_case								=       TransfoXLLMHeadModel(snake_case_	)

      __snake_case								=       load_tf_weights_in_transfo_xl(snake_case_	,				snake_case_	,				snake_case_	)
      # Save pytorch-model
      __snake_case								=       os.path.join(snake_case_	,				snake_case_	)
      __snake_case								=       os.path.join(snake_case_	,				snake_case_	)
      print(f"""Save PyTorch model to {os.path.abspath(snake_case_	)}"""	)
      torch.save(model.state_dict()	,				snake_case_	)
      print(f"""Save configuration file to {os.path.abspath(snake_case_	)}"""	)
      with open(snake_case_	,				'''w'''	,				encoding='''utf-8'''	) as f:
         f.write(config.to_json_string()	)


if __name__ == "__main__":
    snake_case_           =      argparse.ArgumentParser()
    parser.add_argument(
        '--pytorch_dump_folder_path',
        default=None,
        type=str,
        required=True,
        help='Path to the folder to store the PyTorch model or dataset/vocab.',
    )
    parser.add_argument(
        '--tf_checkpoint_path',
        default='',
        type=str,
        help='An optional path to a TensorFlow checkpoint path to be converted.',
    )
    parser.add_argument(
        '--transfo_xl_config_file',
        default='',
        type=str,
        help=(
            'An optional config json file corresponding to the pre-trained BERT model. \n'
            'This specifies the model architecture.'
        ),
    )
    parser.add_argument(
        '--transfo_xl_dataset_file',
        default='',
        type=str,
        help='An optional dataset file to be converted in a vocabulary.',
    )
    snake_case_           =      parser.parse_args()
    convert_transfo_xl_checkpoint_to_pytorch(
        args.tf_checkpoint_path,
        args.transfo_xl_config_file,
        args.pytorch_dump_folder_path,
        args.transfo_xl_dataset_file,
    )

