



"""simple docstring"""

from typing import TYPE_CHECKING

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


lowercase__    :  List[str]							    =	{
    '''configuration_electra''': ['''ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ElectraConfig''', '''ElectraOnnxConfig'''],
    '''tokenization_electra''': ['''ElectraTokenizer'''],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__    :  List[str]							    =	['''ElectraTokenizerFast''']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__    :  Union[str, Any]							    =	[
         '''ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''ElectraForCausalLM''',
         '''ElectraForMaskedLM''',
         '''ElectraForMultipleChoice''',
         '''ElectraForPreTraining''',
         '''ElectraForQuestionAnswering''',
         '''ElectraForSequenceClassification''',
         '''ElectraForTokenClassification''',
         '''ElectraModel''',
         '''ElectraPreTrainedModel''',
         '''load_tf_weights_in_electra''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__    :  List[str]							    =	[
         '''TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFElectraForMaskedLM''',
         '''TFElectraForMultipleChoice''',
         '''TFElectraForPreTraining''',
         '''TFElectraForQuestionAnswering''',
         '''TFElectraForSequenceClassification''',
         '''TFElectraForTokenClassification''',
         '''TFElectraModel''',
         '''TFElectraPreTrainedModel''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     lowercase__    :  Any							    =	[
         '''FlaxElectraForCausalLM''',
         '''FlaxElectraForMaskedLM''',
         '''FlaxElectraForMultipleChoice''',
         '''FlaxElectraForPreTraining''',
         '''FlaxElectraForQuestionAnswering''',
         '''FlaxElectraForSequenceClassification''',
         '''FlaxElectraForTokenClassification''',
         '''FlaxElectraModel''',
         '''FlaxElectraPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
     from .tokenization_electra import ElectraTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_electra_fast import ElectraTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_electra import (
              ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
              ElectraForCausalLM,
              ElectraForMaskedLM,
              ElectraForMultipleChoice,
              ElectraForPreTraining,
              ElectraForQuestionAnswering,
              ElectraForSequenceClassification,
              ElectraForTokenClassification,
              ElectraModel,
              ElectraPreTrainedModel,
              load_tf_weights_in_electra,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_electra import (
              TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFElectraForMaskedLM,
              TFElectraForMultipleChoice,
              TFElectraForPreTraining,
              TFElectraForQuestionAnswering,
              TFElectraForSequenceClassification,
              TFElectraForTokenClassification,
              TFElectraModel,
              TFElectraPreTrainedModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_electra import (
              FlaxElectraForCausalLM,
              FlaxElectraForMaskedLM,
              FlaxElectraForMultipleChoice,
              FlaxElectraForPreTraining,
              FlaxElectraForQuestionAnswering,
              FlaxElectraForSequenceClassification,
              FlaxElectraForTokenClassification,
              FlaxElectraModel,
              FlaxElectraPreTrainedModel,
          )

else:
     import sys

     lowercase__    :  Optional[int]							    =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)










"""simple docstring"""

import argparse
import json

import numpy
import torch

from transformers.models.xlm.tokenization_xlm import VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()




def __lowercase			(							_a		,			_a    ):
    # Load checkpoint
    snake_case_ :     Optional[Any]				=	torch.load(_a		,			map_location='''cpu'''    )

    snake_case_ :     Union[str, Any]				=	chkpt['''model''']

    # We have the base model one level deeper than the original XLM repository
    snake_case_ :     Dict				=	{}
    for k, v in state_dict.items():
        if "pred_layer" in k:
            snake_case_ :     Union[str, Any]				=	v
        else:
            snake_case_ :     Dict				=	v

    snake_case_ :     Union[str, Any]				=	chkpt['''params''']
    snake_case_ :     int				=	{n: v for n, v in config.items() if not isinstance(_a		,			(torch.FloatTensor, numpy.ndarray)    )}

    snake_case_ :     int				=	chkpt['''dico_word2id''']
    snake_case_ :     str				=	{s + '''</w>''' if s.find('''@@'''    ) == -1 and i > 13 else s.replace('''@@'''		,			''''''    ): i for s, i in vocab.items()}

    # Save pytorch-model
    snake_case_ :     Union[str, Any]				=	pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
    snake_case_ :     Union[str, Any]				=	pytorch_dump_folder_path + '''/''' + CONFIG_NAME
    snake_case_ :     Any				=	pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''vocab_file''']

    print(f"Save PyTorch model to {pytorch_weights_dump_path}"    )
    torch.save(_a		,			_a    )

    print(f"Save configuration file to {pytorch_config_dump_path}"    )
    with open(_a		,			'''w'''		,			encoding='''utf-8'''    ) as f:
        f.write(json.dumps(_a		,			indent=2    ) + '''\n'''    )

    print(f"Save vocab file to {pytorch_config_dump_path}"    )
    with open(_a		,			'''w'''		,			encoding='''utf-8'''    ) as f:
        f.write(json.dumps(_a		,			indent=2    ) + '''\n'''    )


if __name__ == "__main__":
     lowercase__    :  Optional[Any]							    =	argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--xlm_checkpoint_path''', default=None, type=str, required=True, help='''Path the official PyTorch dump.'''
     )
     parser.add_argument(
         '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
     )
     lowercase__    :  List[str]							    =	parser.parse_args()
     convert_xlm_checkpoint_to_pytorch(args.xlm_checkpoint_path, args.pytorch_dump_folder_path)





