


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

lowerCAmelCase__      		= {
    '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',
}
lowerCAmelCase__      		= [
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]
def 							_UpperCAmelCase				(UpperCamelCase__     :  List[Any]				,							UpperCamelCase__     :  List[str]				,							UpperCamelCase__     :  str				,							UpperCamelCase__     :  Tuple				,							UpperCamelCase__     :  List[Any]						):
       for attribute in key.split("."						):
              _A      :       Optional[int]				    = getattr(UpperCamelCase__				,							UpperCamelCase__						)

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

       logger.info(f"{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}."						)
def 							_UpperCAmelCase				(UpperCamelCase__     :  str				,							UpperCamelCase__     :  Optional[Any]						):
       _A      :       str				    = []
       _A      :       Optional[int]				    = fairseq_model.state_dict()

       _A      :       List[Any]				    = hf_model.feature_extractor

       # if encoder has different dim to decoder -> use proj_weight
       _A      :       List[Any]				    = None

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

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

       return proj_weight
def 							_UpperCAmelCase				(UpperCamelCase__     :  Optional[Any]				,							UpperCamelCase__     :  Optional[int]				,							UpperCamelCase__     :  Union[str, Any]				,							UpperCamelCase__     :  int				,							UpperCamelCase__     :  Optional[Any]						):
       _A      :       List[Any]				    = full_name.split("conv_layers."						)[-1]
       _A      :       Any				    = name.split("."						)
       _A      :       Union[str, Any]				    = int(items[0]						)
       _A      :       Dict				    = 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."
                     )
                     _A      :       Optional[int]				    = 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."
                     )
                     _A      :       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."
                     )
                     _A      :       List[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."
                     )
                     _A      :       List[Any]				    = value
                     logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."						)
       else:
              unused_weights.append(UpperCamelCase__						)
def 							_UpperCAmelCase				(UpperCamelCase__     :  List[Any]						):
       _A    ,			_A      :       int				    = emb.weight.shape
       _A      :       Any				    = nn.Linear(UpperCamelCase__				,							UpperCamelCase__				,							bias=UpperCamelCase__						)
       _A      :       int				    = emb.weight.data
       return lin_layer
def 							_UpperCAmelCase				(UpperCamelCase__     :  Union[str, Any]						):
       with open(UpperCamelCase__				,							"r"				,							encoding="utf-8"						) as f:
              _A      :       List[Any]				    = f.readlines()
              _A      :       Any				    = [line.split(" "						)[0] for line in lines]

       _A      :       Dict				    = len(UpperCamelCase__						)

       _A      :       Optional[int]				    = {
           "<s>": 0,
           "<pad>": 1,
           "</s>": 2,
           "<unk>": 3,
       }

       vocab_dict.update(dict(zip(UpperCamelCase__				,							range(4				,							num_words + 4						)						)						)						)
       return vocab_dict
@torch.no_grad()
def 							_UpperCAmelCase				(UpperCamelCase__     :  Optional[int]				,							UpperCamelCase__     :  Dict				,							UpperCamelCase__     :  int				,							UpperCamelCase__     :  Any				,							UpperCamelCase__     :  str				,							UpperCamelCase__     :  str				,							UpperCamelCase__     :  int				,							):
       _A      :       int				    = WavaVecaConfig.from_pretrained(UpperCamelCase__						)
       _A      :       Optional[int]				    = SpeechaTextaConfig.from_pretrained(
           UpperCamelCase__				,							vocab_size=UpperCamelCase__				,							decoder_layers=UpperCamelCase__				,							do_stable_layer_norm=UpperCamelCase__						)

       _A      :       Dict				    = WavaVecaFeatureExtractor(
           feature_size=1				,							sampling_rate=16000				,							padding_value=0				,							do_normalize=UpperCamelCase__				,							return_attention_mask=UpperCamelCase__				,							)

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

       # set weights for wav2vec2 encoder
       _A      :       Optional[Any]				    = WavaVecaModel(UpperCamelCase__						)
       _A      :       List[str]				    = recursively_load_weights_wavaveca(model.encoder				,							UpperCamelCase__						)

       _A      :       List[Any]				    = SpeechaTextaForCausalLM(UpperCamelCase__						)
       _A    ,			_A      :       Optional[Any]				    = hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()				,							strict=UpperCamelCase__						)

       # set output linear layer
       unexpected_keys.remove("embed_out"						)
       _A      :       Dict				    = 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}"						)

       _A      :       str				    = SpeechEncoderDecoderModel(encoder=UpperCamelCase__				,							decoder=UpperCamelCase__						)
       _A      :       Dict				    = False

       # add projection layer
       _A      :       Union[str, Any]				    = nn.Parameter(projection_layer.weight						)
       _A      :       str				    = nn.Parameter(projection_layer.bias						)

       _A      :       Optional[Any]				    = create_vocab_dict(UpperCamelCase__						)

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

       _A      :       Optional[int]				    = SpeechaTextaTokenizer(os.path.join(UpperCamelCase__				,							"vocab.json"						)						)
       tokenizer.save_pretrained(UpperCamelCase__						)

       _A      :       Union[str, Any]				    = hf_wavavec.config.to_dict()
       _A      :       Any				    = tokenizer.pad_token_id
       _A      :       Any				    = tokenizer.bos_token_id
       _A      :       Optional[int]				    = tokenizer.eos_token_id
       _A      :       List[Any]				    = "speech_to_text_2"
       _A      :       Tuple				    = "wav2vec2"

       _A      :       Tuple				    = SpeechEncoderDecoderConfig.from_dict(UpperCamelCase__						)

       hf_wavavec.save_pretrained(UpperCamelCase__						)
       feature_extractor.save_pretrained(UpperCamelCase__						)


if __name__ == "__main__":
    lowerCAmelCase__      		= 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_02_24, type=int, help='Vocab size of decoder')
    parser.add_argument('--num_decoder_layers', default=7, type=int, help='Number of decoder layers')

    lowerCAmelCase__      		= 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,
    )







a__:      Dict									=		{0: [2, 3], 1: [0], 2: [1], 3: [4], 4: []}

a__:      str									=		{0: [1, 2, 3], 1: [2], 2: [0], 3: [4], 4: [5], 5: [3]}


def 	UpperCamelCase__(       UpperCamelCase__			:       dict[int, list[int]] ,     UpperCamelCase__			:       int ,     UpperCamelCase__			:       list[bool]  )->list[int]:
       A__											=						True
       A__											=						[]

       for neighbour in graph[vert]:
              if not visited[neighbour]:
                     order += topology_sort(UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__  )

       order.append(UpperCamelCase__  )

       return order


def 	UpperCamelCase__(       UpperCamelCase__			:       dict[int, list[int]] ,     UpperCamelCase__			:       int ,     UpperCamelCase__			:       list[bool]  )->list[int]:
       A__											=						True
       A__											=						[vert]

       for neighbour in reversed_graph[vert]:
              if not visited[neighbour]:
                     component += find_components(UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__  )

       return component


def 	UpperCamelCase__(       UpperCamelCase__			:       dict[int, list[int]]  )->list[list[int]]:
       A__											=						len(UpperCamelCase__  ) * [False]
       A__											=						{vert: [] for vert in range(len(UpperCamelCase__  )  )}

       for vert, neighbours in graph.items():
              for neighbour in neighbours:
                     reversed_graph[neighbour].append(UpperCamelCase__  )

       A__											=						[]
       for i, was_visited in enumerate(UpperCamelCase__  ):
              if not was_visited:
                     order += topology_sort(UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__  )

       A__											=						[]
       A__											=						len(UpperCamelCase__  ) * [False]

       for i in range(len(UpperCamelCase__  )  ):
              A__											=						order[len(UpperCamelCase__  ) - i - 1]
              if not visited[vert]:
                     A__											=						find_components(UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__  )
                     components_list.append(UpperCamelCase__  )

       return components_list





