from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_								=							logging.get_logger(__name__)

UpperCAmelCase_								=							{
    '''facebook/s2t-wav2vec2-large-en-de''': (
        '''https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/config.json'''
    ),
    # See all Speech2Text models at https://huggingface.co/models?filter=speech2text2
}





class      lowerCAmelCase						(					_a			):
     _SCREAMING_SNAKE_CASE     :     Optional[Any]		="""speech_to_text_2"""
     _SCREAMING_SNAKE_CASE     :     str		=["""past_key_values"""]
     _SCREAMING_SNAKE_CASE     :     Tuple		={"""num_attention_heads""": """decoder_attention_heads""", """hidden_size""": """d_model"""}



     def __init__(					self				,      lowerCAmelCase__=10000				,      lowerCAmelCase__=6				,      lowerCAmelCase__=2048				,      lowerCAmelCase__=4				,      lowerCAmelCase__=0.0				,      lowerCAmelCase__=True				,      lowerCAmelCase__="relu"				,      lowerCAmelCase__=256				,      lowerCAmelCase__=0.1				,      lowerCAmelCase__=0.0				,      lowerCAmelCase__=0.0				,      lowerCAmelCase__=0.02				,      lowerCAmelCase__=2				,      lowerCAmelCase__=True				,      lowerCAmelCase__=1				,      lowerCAmelCase__=0				,      lowerCAmelCase__=2				,      lowerCAmelCase__=1024				,      **lowerCAmelCase__				,      ):
            _A=	vocab_size
            _A=	d_model
            _A=	decoder_ffn_dim
            _A=	decoder_layers
            _A=	decoder_attention_heads
            _A=	dropout
            _A=	attention_dropout
            _A=	activation_dropout
            _A=	activation_function
            _A=	init_std
            _A=	decoder_layerdrop
            _A=	use_cache
            _A=	decoder_layers
            _A=	scale_embedding  # scale factor will be sqrt(d_model) if True
            _A=	max_target_positions

            super().__init__(
                pad_token_id=lowerCAmelCase__				,      bos_token_id=lowerCAmelCase__				,      eos_token_id=lowerCAmelCase__				,      decoder_start_token_id=lowerCAmelCase__				,      **lowerCAmelCase__				,      )
import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
      UpperCAmelCase_								=							argparse.ArgumentParser(
          description=(
              '''Extraction some layers of the full BertForMaskedLM or RObertaForMaskedLM for Transfer Learned'''
              ''' Distillation'''
          )
      )
      parser.add_argument('''--model_type''', default='''bert''', choices=['''bert'''])
      parser.add_argument('''--model_name''', default='''bert-base-uncased''', type=str)
      parser.add_argument('''--dump_checkpoint''', default='''serialization_dir/tf_bert-base-uncased_0247911.pth''', type=str)
      parser.add_argument('''--vocab_transform''', action='''store_true''')
      UpperCAmelCase_								=							parser.parse_args()

      if args.model_type == "bert":
            UpperCAmelCase_								=							BertForMaskedLM.from_pretrained(args.model_name)
            UpperCAmelCase_								=							'''bert'''
      else:
            raise ValueError('''args.model_type should be "bert".''')

      UpperCAmelCase_								=							model.state_dict()
      UpperCAmelCase_								=							{}

      for w in ["word_embeddings", "position_embeddings"]:
            UpperCAmelCase_								=							state_dict[F"{prefix}.embeddings.{w}.weight"]
      for w in ["weight", "bias"]:
            UpperCAmelCase_								=							state_dict[F"{prefix}.embeddings.LayerNorm.{w}"]

      UpperCAmelCase_								=							0
      for teacher_idx in [0, 2, 4, 7, 9, 11]:
            for w in ["weight", "bias"]:
                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}"
                  ]
                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}"
                  ]
                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}"
                  ]

                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"
                  ]
                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"
                  ]

                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"
                  ]
                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"
                  ]
                  UpperCAmelCase_								=							state_dict[
                      F"{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"
                  ]
            std_idx += 1

      UpperCAmelCase_								=							state_dict['''cls.predictions.decoder.weight''']
      UpperCAmelCase_								=							state_dict['''cls.predictions.bias''']
      if args.vocab_transform:
            for w in ["weight", "bias"]:
                  UpperCAmelCase_								=							state_dict[F"cls.predictions.transform.dense.{w}"]
                  UpperCAmelCase_								=							state_dict[F"cls.predictions.transform.LayerNorm.{w}"]

      print(F"N layers selected for distillation: {std_idx}")
      print(F"Number of params transferred for distillation: {len(compressed_sd.keys())}")

      print(F"Save transferred checkpoint to {args.dump_checkpoint}.")
      torch.save(compressed_sd, args.dump_checkpoint)