



import argparse

import torch

from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
from transformers.utils import logging


logging.set_verbosity_info()




def     UpperCAmelCase__ (							__snake_case						,				__snake_case						,				__snake_case						,				__snake_case							)   -> Tuple:
  # Initialise PyTorch model
  _A           =    ElectraConfig.from_json_file(__snake_case							)
  print(F'''Building PyTorch model from configuration: {config}'''							)

  if discriminator_or_generator == "discriminator":
    _A           =    ElectraForPreTraining(__snake_case							)
  elif discriminator_or_generator == "generator":
    _A           =    ElectraForMaskedLM(__snake_case							)
  else:
    raise ValueError('''The discriminator_or_generator argument should be either \'discriminator\' or \'generator\''''							)

  # Load weights from tf checkpoint
  load_tf_weights_in_electra(
      __snake_case						,				__snake_case						,				__snake_case						,				discriminator_or_generator=__snake_case							)

  # Save pytorch-model
  print(F'''Save PyTorch model to {pytorch_dump_path}'''							)
  torch.save(model.state_dict()						,				__snake_case							)


if __name__ == "__main__":
    __lowerCamelCase  				=					argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
    )
    parser.add_argument(
        """--config_file""",
        default=None,
        type=str,
        required=True,
        help="""The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.""",
    )
    parser.add_argument(
        """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
    )
    parser.add_argument(
        """--discriminator_or_generator""",
        default=None,
        type=str,
        required=True,
        help=(
            """Whether to export the generator or the discriminator. Should be a string, either 'discriminator' or """
            """'generator'."""
        ),
    )
    __lowerCamelCase  				=					parser.parse_args()
    convert_tf_checkpoint_to_pytorch(
        args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
    )




import copy

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


__lowerCamelCase  				=					logging.get_logger(__name__)

__lowerCamelCase  				=					{}



class 		_snake_case      (			lowerCamelCase				):





 """simple docstring"""



 lowerCamelCase_          =			'''fsmt_decoder'''
 def __init__(						self	,    a=0	,    a=0					)	->							str:




   """simple docstring"""







   super().__init__()
   _A           =    vocab_size
   _A           =    bos_token_id







class 		_snake_case      (			lowerCamelCase				):





 """simple docstring"""



 lowerCamelCase_          =			'''fsmt'''
 lowerCamelCase_          =			{'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}
 def __init__(						self	,    a=["en", "de"]	,    a=4_2_0_2_4	,    a=4_2_0_2_4	,    a="relu"	,    a=1_0_2_4	,    a=2_0_0	,    a=1_0_2_4	,    a=4_0_9_6	,    a=1_2	,    a=1_6	,    a=0.0	,    a=4_0_9_6	,    a=1_2	,    a=1_6	,    a=0.0	,    a=0.0	,    a=0.1	,    a=0.0	,    a=0.02	,    a=2	,    a=True	,    a=True	,    a=False	,    a=5	,    a=1.0	,    a=False	,    a=True	,    a=1	,    a=0	,    a=2	,    a=2	,    **a	,    )	->							Union[str, Any]:




   """simple docstring"""







   _A           =    langs
   _A           =    src_vocab_size
   _A           =    tgt_vocab_size
   _A           =    d_model  # encoder_embed_dim and decoder_embed_dim

   _A           =    encoder_ffn_dim
   _A           =    _A           =    encoder_layers
   _A           =    encoder_attention_heads
   _A           =    encoder_layerdrop
   _A           =    decoder_layerdrop
   _A           =    decoder_ffn_dim
   _A           =    decoder_layers
   _A           =    decoder_attention_heads
   _A           =    max_position_embeddings
   _A           =    init_std  # Normal(0, this parameter)
   _A           =    activation_function

   _A           =    DecoderConfig(vocab_size=a	,    bos_token_id=a					)
   if "decoder" in common_kwargs:
     del common_kwargs["decoder"]

   _A           =    scale_embedding  # scale factor will be sqrt(d_model) if True

   # 3 Types of Dropout
   _A           =    attention_dropout
   _A           =    activation_dropout
   _A           =    dropout

   _A           =    use_cache
   super().__init__(
       pad_token_id=a	,    bos_token_id=a	,    eos_token_id=a	,    decoder_start_token_id=a	,    is_encoder_decoder=a	,    tie_word_embeddings=a	,    forced_eos_token_id=a	,    max_length=a	,    num_beams=a	,    length_penalty=a	,    early_stopping=a	,    **a	,    )



 def lowercase_		(						self					)	->							Any:




   """simple docstring"""







   _A           =    copy.deepcopy(self.__dict__					)
   _A           =    self.decoder.to_dict()
   _A           =    self.__class__.model_type
   return output