

'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


lowerCamelCase :      Optional[Any]					 =     {
    "configuration_roberta_prelayernorm": [
        "ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "RobertaPreLayerNormConfig",
        "RobertaPreLayerNormOnnxConfig",
    ],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCamelCase :      Optional[int]					 =     [
         "ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
         "RobertaPreLayerNormForCausalLM",
         "RobertaPreLayerNormForMaskedLM",
         "RobertaPreLayerNormForMultipleChoice",
         "RobertaPreLayerNormForQuestionAnswering",
         "RobertaPreLayerNormForSequenceClassification",
         "RobertaPreLayerNormForTokenClassification",
         "RobertaPreLayerNormModel",
         "RobertaPreLayerNormPreTrainedModel",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCamelCase :      List[str]					 =     [
         "TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
         "TFRobertaPreLayerNormForCausalLM",
         "TFRobertaPreLayerNormForMaskedLM",
         "TFRobertaPreLayerNormForMultipleChoice",
         "TFRobertaPreLayerNormForQuestionAnswering",
         "TFRobertaPreLayerNormForSequenceClassification",
         "TFRobertaPreLayerNormForTokenClassification",
         "TFRobertaPreLayerNormMainLayer",
         "TFRobertaPreLayerNormModel",
         "TFRobertaPreLayerNormPreTrainedModel",
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowerCamelCase :      int					 =     [
         "FlaxRobertaPreLayerNormForCausalLM",
         "FlaxRobertaPreLayerNormForMaskedLM",
         "FlaxRobertaPreLayerNormForMultipleChoice",
         "FlaxRobertaPreLayerNormForQuestionAnswering",
         "FlaxRobertaPreLayerNormForSequenceClassification",
         "FlaxRobertaPreLayerNormForTokenClassification",
         "FlaxRobertaPreLayerNormModel",
         "FlaxRobertaPreLayerNormPreTrainedModel",
     ]


if TYPE_CHECKING:
     from .configuration_roberta_prelayernorm import (
         ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP,
         RobertaPreLayerNormConfig,
         RobertaPreLayerNormOnnxConfig,
     )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_roberta_prelayernorm import (
              ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
              RobertaPreLayerNormForCausalLM,
              RobertaPreLayerNormForMaskedLM,
              RobertaPreLayerNormForMultipleChoice,
              RobertaPreLayerNormForQuestionAnswering,
              RobertaPreLayerNormForSequenceClassification,
              RobertaPreLayerNormForTokenClassification,
              RobertaPreLayerNormModel,
              RobertaPreLayerNormPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_roberta_prelayernorm import (
              TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFRobertaPreLayerNormForCausalLM,
              TFRobertaPreLayerNormForMaskedLM,
              TFRobertaPreLayerNormForMultipleChoice,
              TFRobertaPreLayerNormForQuestionAnswering,
              TFRobertaPreLayerNormForSequenceClassification,
              TFRobertaPreLayerNormForTokenClassification,
              TFRobertaPreLayerNormMainLayer,
              TFRobertaPreLayerNormModel,
              TFRobertaPreLayerNormPreTrainedModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_roberta_prelayernorm import (
              FlaxRobertaPreLayerNormForCausalLM,
              FlaxRobertaPreLayerNormForMaskedLM,
              FlaxRobertaPreLayerNormForMultipleChoice,
              FlaxRobertaPreLayerNormForQuestionAnswering,
              FlaxRobertaPreLayerNormForSequenceClassification,
              FlaxRobertaPreLayerNormForTokenClassification,
              FlaxRobertaPreLayerNormModel,
              FlaxRobertaPreLayerNormPreTrainedModel,
          )

else:
     import sys

     lowerCamelCase :      int					 =     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


"""simple docstring"""



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)



