


'''simple docstring'''
import argparse
import json
import os

import fairseq
import torch
from torch import nn

from transformers import (
    SpeechaTextaConfig,
    SpeechaTextaForCausalLM,
    SpeechaTextaTokenizer,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


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

UpperCamelCase__    					=      {
    '''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''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
}
UpperCamelCase__    					=      [
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	int:
    for attribute in key.split('''.'''						):
        UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					lowerCAmelCase__						)

    if weight_type is not None:
        UpperCAmelCase__			:		Optional[Any]													=				getattr(lowerCAmelCase__  ,					lowerCAmelCase__						).shape
    else:
        UpperCAmelCase__			:		List[Any]													=				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":
        UpperCAmelCase__			:		int													=				value
    elif weight_type == "weight_g":
        UpperCAmelCase__			:		Union[str, Any]													=				value
    elif weight_type == "weight_v":
        UpperCAmelCase__			:		Dict													=				value
    elif weight_type == "bias":
        UpperCAmelCase__			:		Union[str, Any]													=				value
    else:
        UpperCAmelCase__			:		int													=				value

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





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Optional[int]:
    UpperCAmelCase__			:		int													=				[]
    UpperCAmelCase__			:		List[Any]													=				fairseq_model.state_dict()

    UpperCAmelCase__			:		Union[str, Any]													=				hf_model.feature_extractor

    # if encoder has different dim to decoder -> use proj_weight
    UpperCAmelCase__			:		str													=				None

    for name, value in fairseq_dict.items():
        UpperCAmelCase__			:		List[Any]													=				False
        if "conv_layers" in name:
            load_conv_layer(
                lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					hf_model.config.feat_extract_norm == '''group'''  ,					)
            UpperCAmelCase__			:		Tuple													=				True
        elif name.split('''.'''						)[0] == "proj":
            UpperCAmelCase__			:		int													=				fairseq_model.proj
            UpperCAmelCase__			:		int													=				True
        else:
            for key, mapped_key in MAPPING.items():
                if key in name or key.split('''w2v_model.'''						)[-1] == name.split('''.'''						)[0]:
                    UpperCAmelCase__			:		Optional[int]													=				True
                    if "*" in mapped_key:
                        UpperCAmelCase__			:		Optional[Any]													=				name.split(lowerCAmelCase__						)[0].split('''.'''						)[-2]
                        UpperCAmelCase__			:		Union[str, Any]													=				mapped_key.replace('''*'''  ,					lowerCAmelCase__						)
                    if "weight_g" in name:
                        UpperCAmelCase__			:		List[Any]													=				'''weight_g'''
                    elif "weight_v" in name:
                        UpperCAmelCase__			:		Dict													=				'''weight_v'''
                    elif "bias" in name:
                        UpperCAmelCase__			:		int													=				'''bias'''
                    elif "weight" in name:
                        UpperCAmelCase__			:		Optional[Any]													=				'''weight'''
                    else:
                        UpperCAmelCase__			:		List[Any]													=				None
                    set_recursively(lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)
                continue
        if not is_used:
            unused_weights.append(lowerCAmelCase__						)

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

    return proj_weight





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	List[str]:
    UpperCAmelCase__			:		Union[str, Any]													=				full_name.split('''conv_layers.'''						)[-1]
    UpperCAmelCase__			:		int													=				name.split('''.'''						)
    UpperCAmelCase__			:		str													=				int(items[0]						)
    UpperCAmelCase__			:		Tuple													=				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."""
            )
            UpperCAmelCase__			:		List[str]													=				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."""
            )
            UpperCAmelCase__			:		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:
            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."
            )
            UpperCAmelCase__			:		str													=				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."""
            )
            UpperCAmelCase__			:		List[Any]													=				value
            logger.info(F"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""						)
    else:
        unused_weights.append(lowerCAmelCase__						)





def 					a__       (					lowerCAmelCase__						)    ->	int:
    UpperCAmelCase__						,						UpperCAmelCase__			:		Any													=				emb.weight.shape
    UpperCAmelCase__			:		Optional[int]													=				nn.Linear(lowerCAmelCase__  ,					lowerCAmelCase__  ,					bias=lowerCAmelCase__						)
    UpperCAmelCase__			:		Any													=				emb.weight.data
    return lin_layer





def 					a__       (					lowerCAmelCase__						)    ->	List[Any]:
    with open(lowerCAmelCase__  ,					'''r'''  ,					encoding='''utf-8'''						) as f:
        UpperCAmelCase__			:		Tuple													=				f.readlines()
        UpperCAmelCase__			:		str													=				[line.split(''' '''						)[0] for line in lines]

    UpperCAmelCase__			:		List[Any]													=				len(lowerCAmelCase__						)

    UpperCAmelCase__			:		List[str]													=				{
        '''<s>''': 0,
        '''<pad>''': 1,
        '''</s>''': 2,
        '''<unk>''': 3,
    }

    vocab_dict.update(dict(zip(lowerCAmelCase__  ,					range(4  ,					num_words + 4						)						)						)						)
    return vocab_dict





@torch.no_grad()
def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__  ,					)    ->	Union[str, Any]:
    UpperCAmelCase__			:		str													=				WavaVecaConfig.from_pretrained(lowerCAmelCase__						)
    UpperCAmelCase__			:		Optional[int]													=				SpeechaTextaConfig.from_pretrained(
        lowerCAmelCase__  ,					vocab_size=lowerCAmelCase__  ,					decoder_layers=lowerCAmelCase__  ,					do_stable_layer_norm=lowerCAmelCase__						)

    UpperCAmelCase__			:		Tuple													=				WavaVecaFeatureExtractor(
        feature_size=1  ,					sampling_rate=1_60_00  ,					padding_value=0  ,					do_normalize=lowerCAmelCase__  ,					return_attention_mask=lowerCAmelCase__  ,					)

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

    # set weights for wav2vec2 encoder
    UpperCAmelCase__			:		Optional[Any]													=				WavaVecaModel(lowerCAmelCase__						)
    UpperCAmelCase__			:		Any													=				recursively_load_weights_wavaveca(model.encoder  ,					lowerCAmelCase__						)

    UpperCAmelCase__			:		List[str]													=				SpeechaTextaForCausalLM(lowerCAmelCase__						)
    UpperCAmelCase__						,						UpperCAmelCase__			:		Union[str, Any]													=				hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()  ,					strict=lowerCAmelCase__						)

    # set output linear layer
    unexpected_keys.remove('''embed_out'''						)
    UpperCAmelCase__			:		str													=				nn.Parameter(model.decoder.embed_out.detach()						)

    # layer norm is init to identity matrix so leaving it is fine
    logger.warning(F"""The following keys are missing when loading the decoder weights: {missing_keys}"""						)
    logger.warning(F"""The following keys are unexpected when loading the decoder weights: {unexpected_keys}"""						)

    UpperCAmelCase__			:		Any													=				SpeechEncoderDecoderModel(encoder=lowerCAmelCase__  ,					decoder=lowerCAmelCase__						)
    UpperCAmelCase__			:		List[Any]													=				False

    # add projection layer
    UpperCAmelCase__			:		List[str]													=				nn.Parameter(projection_layer.weight						)
    UpperCAmelCase__			:		Any													=				nn.Parameter(projection_layer.bias						)

    UpperCAmelCase__			:		List[str]													=				create_vocab_dict(lowerCAmelCase__						)

    with open(os.path.join(lowerCAmelCase__  ,					'''vocab.json'''						)  ,					'''w'''						) as fp:
        json.dump(lowerCAmelCase__  ,					lowerCAmelCase__						)

    UpperCAmelCase__			:		List[str]													=				SpeechaTextaTokenizer(os.path.join(lowerCAmelCase__  ,					'''vocab.json'''						)						)
    tokenizer.save_pretrained(lowerCAmelCase__						)

    UpperCAmelCase__			:		Dict													=				hf_wavavec.config.to_dict()
    UpperCAmelCase__			:		Any													=				tokenizer.pad_token_id
    UpperCAmelCase__			:		List[Any]													=				tokenizer.bos_token_id
    UpperCAmelCase__			:		Dict													=				tokenizer.eos_token_id
    UpperCAmelCase__			:		Tuple													=				'''speech_to_text_2'''
    UpperCAmelCase__			:		Any													=				'''wav2vec2'''

    UpperCAmelCase__			:		Tuple													=				SpeechEncoderDecoderConfig.from_dict(lowerCAmelCase__						)

    hf_wavavec.save_pretrained(lowerCAmelCase__						)
    feature_extractor.save_pretrained(lowerCAmelCase__						)


if __name__ == "__main__":
  UpperCamelCase__    					=      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(
      '''--encoder_config_path''',
      default='''facebook/wav2vec2-large-lv60''',
      type=str,
      help='''Path to hf encoder wav2vec2 checkpoint config''',
  )
  parser.add_argument(
      '''--decoder_config_path''',
      default='''facebook/s2t-small-mustc-en-fr-st''',
      type=str,
      help='''Path to hf decoder s2t checkpoint config''',
  )
  parser.add_argument('''--vocab_size''', default=1_0_2_2_4, type=int, help='''Vocab size of decoder''')
  parser.add_argument('''--num_decoder_layers''', default=7, type=int, help='''Number of decoder layers''')

  UpperCamelCase__    					=      parser.parse_args()
  convert_wavaveca_checkpoint(
      args.checkpoint_path,
      args.pytorch_dump_folder_path,
      args.dict_path,
      encoder_config_path=args.encoder_config_path,
      decoder_config_path=args.decoder_config_path,
      vocab_size=args.vocab_size,
      num_decoder_layers=args.num_decoder_layers,
  )




'''simple docstring'''
import argparse
import shlex

import runhouse as rh


if __name__ == "__main__":
  # Refer to https://runhouse-docs.readthedocs-hosted.com/en/latest/api/python/cluster.html#hardware-setup for cloud access
  # setup instructions, if using on-demand hardware

  # If user passes --user <user> --host <host> --key_path <key_path> <example> <args>, fill them in as BYO cluster
  # If user passes --instance <instance> --provider <provider> <example> <args>, fill them in as on-demand cluster
  # Throw an error if user passes both BYO and on-demand cluster args
  # Otherwise, use default values
  UpperCamelCase__    					=      argparse.ArgumentParser()
  parser.add_argument('''--user''', type=str, default='''ubuntu''')
  parser.add_argument('''--host''', type=str, default='''localhost''')
  parser.add_argument('''--key_path''', type=str, default=None)
  parser.add_argument('''--instance''', type=str, default='''V100:1''')
  parser.add_argument('''--provider''', type=str, default='''cheapest''')
  parser.add_argument('''--use_spot''', type=bool, default=False)
  parser.add_argument('''--example''', type=str, default='''pytorch/text-generation/run_generation.py''')
  UpperCamelCase__					,   UpperCamelCase__    					=      parser.parse_known_args()
  if args.host != "localhost":
    if args.instance != "V100:1" or args.provider != "cheapest":
      raise ValueError('''Cannot specify both BYO and on-demand cluster args''')
    UpperCamelCase__    					=      rh.cluster(
        name='''rh-cluster''', ips=[args.host], ssh_creds={'''ssh_user''': args.user, '''ssh_private_key''': args.key_path}
    )
  else:
    UpperCamelCase__    					=      rh.cluster(
        name='''rh-cluster''', instance_type=args.instance, provider=args.provider, use_spot=args.use_spot
    )
  UpperCamelCase__    					=      args.example.rsplit('''/''', 1)[0]

  # Set up remote environment
  cluster.install_packages(['''pip:./'''])  # Installs transformers from local source
  # Note transformers is copied into the home directory on the remote machine, so we can install from there
  cluster.run([F"""pip install -r transformers/examples/{example_dir}/requirements.txt"""])
  cluster.run(['''pip install torch --upgrade --extra-index-url https://download.pytorch.org/whl/cu117'''])

  # Run example. You can bypass the CLI wrapper and paste your own code here.
  cluster.run([F"""python transformers/examples/{args.example} {' '.join(shlex.quote(arg) for arg in unknown)}"""])

  # Alternatively, we can just import and run a training function (especially if there's no wrapper CLI):
  # from my_script... import train
  # reqs = ['pip:./', 'torch', 'datasets', 'accelerate', 'evaluate', 'tqdm', 'scipy', 'scikit-learn', 'tensorboard']
  # launch_train_gpu = rh.function(fn=train,
  #                                system=gpu,
  #                                reqs=reqs,
  #                                name='train_bert_glue')
  #
  # We can pass in arguments just like we would to a function:
  # launch_train_gpu(num_epochs = 3, lr = 2e-5, seed = 42, batch_size = 16
  #                  stream_logs=True)
