
"""simple docstring"""



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


UpperCAmelCase__          = logging.get_logger(__name__)

UpperCAmelCase__          = {
    'transfo-xl-wt103': 'https://huggingface.co/transfo-xl-wt103/resolve/main/config.json',
}

class       lowerCAmelCase__				(						A_	):
     __a						=					"""transfo-xl"""
     __a						=					["""mems"""]
     __a						=					{
         """n_token""": """vocab_size""",
         """hidden_size""": """d_model""",
         """num_attention_heads""": """n_head""",
         """num_hidden_layers""": """n_layer""",
     }







     def __init__(							self			:			List[Any]    ,      _lowerCamelCase			:			int=267735    ,      _lowerCamelCase			:			Dict=[20000, 40000, 200000]    ,      _lowerCamelCase			:			Optional[Any]=1024    ,      _lowerCamelCase			:			Dict=1024    ,      _lowerCamelCase			:			List[str]=16    ,      _lowerCamelCase			:			Any=64    ,      _lowerCamelCase			:			Optional[int]=4096    ,      _lowerCamelCase			:			Dict=4    ,      _lowerCamelCase			:			str=False    ,      _lowerCamelCase			:			str=18    ,      _lowerCamelCase			:			Dict=1600    ,      _lowerCamelCase			:			Union[str, Any]=1000    ,      _lowerCamelCase			:			Dict=True    ,      _lowerCamelCase			:			int=True    ,      _lowerCamelCase			:			List[str]=0    ,      _lowerCamelCase			:			Union[str, Any]=-1    ,      _lowerCamelCase			:			int=True    ,      _lowerCamelCase			:			Tuple=0.1    ,      _lowerCamelCase			:			str=0.0    ,      _lowerCamelCase			:			Optional[Any]=True    ,      _lowerCamelCase			:			Dict="normal"    ,      _lowerCamelCase			:			List[Any]=0.0_1    ,      _lowerCamelCase			:			Union[str, Any]=0.0_1    ,      _lowerCamelCase			:			Optional[int]=0.0_2    ,      _lowerCamelCase			:			List[str]=1e-5    ,      _lowerCamelCase			:			Dict=0    ,      **_lowerCamelCase			:			Optional[Any]    ,      ):
       _snake_case					     =       vocab_size
       _snake_case					     =       []
       self.cutoffs.extend(_lowerCamelCase			)
       if proj_share_all_but_first:
         _snake_case					     =       [False] + [True] * len(self.cutoffs			)
       else:
         _snake_case					     =       [False] + [False] * len(self.cutoffs			)
       _snake_case					     =       d_model
       _snake_case					     =       d_embed
       _snake_case					     =       d_head
       _snake_case					     =       d_inner
       _snake_case					     =       div_val
       _snake_case					     =       pre_lnorm
       _snake_case					     =       n_layer
       _snake_case					     =       n_head
       _snake_case					     =       mem_len
       _snake_case					     =       same_length
       _snake_case					     =       attn_type
       _snake_case					     =       clamp_len
       _snake_case					     =       sample_softmax
       _snake_case					     =       adaptive
       _snake_case					     =       dropout
       _snake_case					     =       dropatt
       _snake_case					     =       untie_r
       _snake_case					     =       init
       _snake_case					     =       init_range
       _snake_case					     =       proj_init_std
       _snake_case					     =       init_std
       _snake_case					     =       layer_norm_epsilon
       super().__init__(eos_token_id=_lowerCamelCase    ,      **_lowerCamelCase			)







     @property
     def      lowercase	(							self			:			Any			):
       # Message copied from Transformer-XL documentation
       logger.info(f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''			)
       return -1







     @max_position_embeddings.setter
     def      lowercase	(							self			:			Optional[int]    ,      _lowerCamelCase			:			Optional[Any]			):
       # Message copied from Transformer-XL documentation
       raise NotImplementedError(
           f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''			)





"""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)



