





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()
a__:      str									=		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',
    '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',
}
a__:      str									=		[
    '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]  )->str:
       for attribute in key.split('''.'''  ):
              A__											=						getattr(UpperCamelCase__ ,     UpperCamelCase__  )

       if weight_type is not None:
              A__											=						getattr(UpperCamelCase__ ,     UpperCamelCase__  ).shape
       else:
              A__											=						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__											=						value
       elif weight_type == "weight_g":
              A__											=						value
       elif weight_type == "weight_v":
              A__											=						value
       elif weight_type == "bias":
              A__											=						value
       else:
              A__											=						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]  )->Optional[Any]:
       A__											=						[]
       A__											=						fairseq_model.state_dict()

       A__											=						hf_model.feature_extractor

       # if encoder has different dim to decoder -> use proj_weight
       A__											=						None

       for name, value in fairseq_dict.items():
              A__											=						False
              if "conv_layers" in name:
                     load_conv_layer(
                         UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__ ,     hf_model.config.feat_extract_norm == '''group''' ,     )
                     A__											=						True
              elif name.split('''.'''  )[0] == "proj":
                     A__											=						fairseq_model.proj
                     A__											=						True
              else:
                     for key, mapped_key in MAPPING.items():
                            if key in name or key.split('''w2v_model.'''  )[-1] == name.split('''.'''  )[0]:
                                   A__											=						True
                                   if "*" in mapped_key:
                                          A__											=						name.split(UpperCamelCase__  )[0].split('''.'''  )[-2]
                                          A__											=						mapped_key.replace('''*''' ,     UpperCamelCase__  )
                                   if "weight_g" in name:
                                          A__											=						'''weight_g'''
                                   elif "weight_v" in name:
                                          A__											=						'''weight_v'''
                                   elif "bias" in name:
                                          A__											=						'''bias'''
                                   elif "weight" in name:
                                          A__											=						'''weight'''
                                   else:
                                          A__											=						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]  )->Any:
       A__											=						full_name.split('''conv_layers.'''  )[-1]
       A__											=						name.split('''.'''  )
       A__											=						int(items[0]  )
       A__											=						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__											=						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__											=						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__											=						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__											=						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]  )->int:
       A__					,	A__											=						emb.weight.shape
       A__											=						nn.Linear(UpperCamelCase__ ,     UpperCamelCase__ ,     bias=UpperCamelCase__  )
       A__											=						emb.weight.data
       return lin_layer


def 	UpperCamelCase__(       UpperCamelCase__			:       Union[str, Any]  )->Union[str, Any]:
       with open(UpperCamelCase__ ,     '''r''' ,     encoding='''utf-8'''  ) as f:
              A__											=						f.readlines()
              A__											=						[line.split(''' '''  )[0] for line in lines]

       A__											=						len(UpperCamelCase__  )

       A__											=						{
           '''<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 ,     )->Tuple:
       A__											=						WavaVecaConfig.from_pretrained(UpperCamelCase__  )
       A__											=						SpeechaTextaConfig.from_pretrained(
           UpperCamelCase__ ,     vocab_size=UpperCamelCase__ ,     decoder_layers=UpperCamelCase__ ,     do_stable_layer_norm=UpperCamelCase__  )

       A__											=						WavaVecaFeatureExtractor(
           feature_size=1 ,     sampling_rate=1_60_00 ,     padding_value=0 ,     do_normalize=UpperCamelCase__ ,     return_attention_mask=UpperCamelCase__ ,     )

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

       # set weights for wav2vec2 encoder
       A__											=						WavaVecaModel(UpperCamelCase__  )
       A__											=						recursively_load_weights_wavaveca(model.encoder ,     UpperCamelCase__  )

       A__											=						SpeechaTextaForCausalLM(UpperCamelCase__  )
       A__					,	A__											=						hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() ,     strict=UpperCamelCase__  )

       # set output linear layer
       unexpected_keys.remove('''embed_out'''  )
       A__											=						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__											=						SpeechEncoderDecoderModel(encoder=UpperCamelCase__ ,     decoder=UpperCamelCase__  )
       A__											=						False

       # add projection layer
       A__											=						nn.Parameter(projection_layer.weight  )
       A__											=						nn.Parameter(projection_layer.bias  )

       A__											=						create_vocab_dict(UpperCamelCase__  )

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

       A__											=						SpeechaTextaTokenizer(os.path.join(UpperCamelCase__ ,     '''vocab.json'''  )  )
       tokenizer.save_pretrained(UpperCamelCase__  )

       A__											=						hf_wavavec.config.to_dict()
       A__											=						tokenizer.pad_token_id
       A__											=						tokenizer.bos_token_id
       A__											=						tokenizer.eos_token_id
       A__											=						'''speech_to_text_2'''
       A__											=						'''wav2vec2'''

       A__											=						SpeechEncoderDecoderConfig.from_dict(UpperCamelCase__  )

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


if __name__ == "__main__":
       a__:      Union[str, Any]									=		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=10_224, type=int, help='Vocab size of decoder')
       parser.add_argument('--num_decoder_layers', default=7, type=int, help='Number of decoder layers')

       a__:      Any									=		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,
       )












import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


a__:      int									=		logging.getLogger(__name__)




class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase=None,__lowerCamelCase=None ):
        A__											=						self.layer[current_layer](__lowerCamelCase,__lowerCamelCase,head_mask[current_layer] )

        A__											=						layer_outputs[0]

        return hidden_states




@add_start_docstrings(
    '''The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''				,					UpperCamelCase__				,					)
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):



 def __init__(							self,__lowerCamelCase ):
        super().__init__(__lowerCamelCase )

        A__											=						BertEncoderWithPabee(__lowerCamelCase )

        self.init_weights()
        A__											=						0
        A__											=						0
        A__											=						0

        A__											=						0



 def 				UpperCamelCase			(							self,__lowerCamelCase ):
        A__											=						threshold



 def 				UpperCamelCase			(							self,__lowerCamelCase ):
        A__											=						patience



 def 				UpperCamelCase			(							self ):
        A__											=						0
        A__											=						0



 def 				UpperCamelCase			(							self ):
        A__											=						self.inference_layers_num / self.inference_instances_num
        A__											=						(
            f"*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ="
            f" {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***"
        )
        print(__lowerCamelCase )




 @add_start_docstrings_to_model_forward(__lowerCamelCase )
 def 				UpperCamelCase			(							self,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=False,):

        if input_ids is not None and inputs_embeds is not None:
               raise ValueError('''You cannot specify both input_ids and inputs_embeds at the same time''' )
        elif input_ids is not None:
               A__											=						input_ids.size()
        elif inputs_embeds is not None:
               A__											=						inputs_embeds.size()[:-1]
        else:
               raise ValueError('''You have to specify either input_ids or inputs_embeds''' )

        A__											=						input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
               A__											=						torch.ones(__lowerCamelCase,device=__lowerCamelCase )
        if token_type_ids is None:
               A__											=						torch.zeros(__lowerCamelCase,dtype=torch.long,device=__lowerCamelCase )

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        A__											=						self.get_extended_attention_mask(__lowerCamelCase,__lowerCamelCase,__lowerCamelCase )

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
               A__					,	A__					,	A__											=						encoder_hidden_states.size()
               A__											=						(encoder_batch_size, encoder_sequence_length)
               if encoder_attention_mask is None:
                      A__											=						torch.ones(__lowerCamelCase,device=__lowerCamelCase )
               A__											=						self.invert_attention_mask(__lowerCamelCase )
        else:
               A__											=						None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        A__											=						self.get_head_mask(__lowerCamelCase,self.config.num_hidden_layers )

        A__											=						self.embeddings(
            input_ids=__lowerCamelCase,position_ids=__lowerCamelCase,token_type_ids=__lowerCamelCase,inputs_embeds=__lowerCamelCase )
        A__											=						embedding_output

        if self.training:
               A__											=						[]
               for i in range(self.config.num_hidden_layers ):
                      A__											=						self.encoder.adaptive_forward(
                          __lowerCamelCase,current_layer=__lowerCamelCase,attention_mask=__lowerCamelCase,head_mask=__lowerCamelCase )

                      A__											=						self.pooler(__lowerCamelCase )
                      A__											=						output_layers[i](output_dropout(__lowerCamelCase ) )
                      res.append(__lowerCamelCase )
        elif self.patience == 0:  # Use all layers for inference
               A__											=						self.encoder(
                   __lowerCamelCase,attention_mask=__lowerCamelCase,head_mask=__lowerCamelCase,encoder_hidden_states=__lowerCamelCase,encoder_attention_mask=__lowerCamelCase,)
               A__											=						self.pooler(encoder_outputs[0] )
               A__											=						[output_layers[self.config.num_hidden_layers - 1](__lowerCamelCase )]
        else:
               A__											=						0
               A__											=						None
               A__											=						0
               for i in range(self.config.num_hidden_layers ):
                      calculated_layer_num += 1
                      A__											=						self.encoder.adaptive_forward(
                          __lowerCamelCase,current_layer=__lowerCamelCase,attention_mask=__lowerCamelCase,head_mask=__lowerCamelCase )

                      A__											=						self.pooler(__lowerCamelCase )
                      A__											=						output_layers[i](__lowerCamelCase )
                      if regression:
                             A__											=						logits.detach()
                             if patient_result is not None:
                                    A__											=						patient_result.detach()
                             if (patient_result is not None) and torch.abs(patient_result - labels ) < self.regression_threshold:
                                    patient_counter += 1
                             else:
                                    A__											=						0
                      else:
                             A__											=						logits.detach().argmax(dim=1 )
                             if patient_result is not None:
                                    A__											=						patient_result.detach().argmax(dim=1 )
                             if (patient_result is not None) and torch.all(labels.eq(__lowerCamelCase ) ):
                                    patient_counter += 1
                             else:
                                    A__											=						0

                      A__											=						logits
                      if patient_counter == self.patience:
                             break
               A__											=						[patient_result]
               self.inference_layers_num += calculated_layer_num
               self.inference_instances_num += 1

        return res




@add_start_docstrings(
    '''Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. '''				,					UpperCamelCase__				,					)
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):



 def __init__(							self,__lowerCamelCase ):
        super().__init__(__lowerCamelCase )
        A__											=						config.num_labels

        A__											=						BertModelWithPabee(__lowerCamelCase )
        A__											=						nn.Dropout(config.hidden_dropout_prob )
        A__											=						nn.ModuleList(
            [nn.Linear(config.hidden_size,self.config.num_labels ) for _ in range(config.num_hidden_layers )] )

        self.init_weights()




 @add_start_docstrings_to_model_forward(__lowerCamelCase )
 def 				UpperCamelCase			(							self,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,):
        A__											=						self.bert(
            input_ids=__lowerCamelCase,attention_mask=__lowerCamelCase,token_type_ids=__lowerCamelCase,position_ids=__lowerCamelCase,head_mask=__lowerCamelCase,inputs_embeds=__lowerCamelCase,output_dropout=self.dropout,output_layers=self.classifiers,regression=self.num_labels == 1,)

        A__											=						(logits[-1],)

        if labels is not None:
               A__											=						None
               A__											=						0
               for ix, logits_item in enumerate(__lowerCamelCase ):
                      if self.num_labels == 1:
                             #  We are doing regression
                             A__											=						MSELoss()
                             A__											=						loss_fct(logits_item.view(-1 ),labels.view(-1 ) )
                      else:
                             A__											=						CrossEntropyLoss()
                             A__											=						loss_fct(logits_item.view(-1,self.num_labels ),labels.view(-1 ) )
                      if total_loss is None:
                             A__											=						loss
                      else:
                             total_loss += loss * (ix + 1)
                      total_weights += ix + 1
               A__											=						(total_loss / total_weights,) + outputs

        return outputs





