

'''simple docstring'''
import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


lowerCamelCase			:				List[str] 			=					{
    "cola": 2,
    "mnli": 3,
    "mrpc": 2,
    "sst-2": 2,
    "sts-b": 1,
    "qqp": 2,
    "qnli": 2,
    "rte": 2,
    "wnli": 2,
}


logging.set_verbosity_info()


def 	_lowerCAmelCase (						_UpperCamelCase						:					Optional[Any]							,			_UpperCamelCase						:					List[Any]							,			_UpperCamelCase						:					Dict							,			_UpperCamelCase						:					Any=None    )				->      Union[str, Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   =XLNetConfig.from_json_file(_UpperCamelCase    )

     _SCREAMING_SNAKE_CASE							   =finetuning_task.lower() if finetuning_task is not None else ''
     if finetuning_task in GLUE_TASKS_NUM_LABELS:
          print(f"Building PyTorch XLNetForSequenceClassification model from configuration: {config}"    )
          _SCREAMING_SNAKE_CASE							   =finetuning_task
          _SCREAMING_SNAKE_CASE							   =GLUE_TASKS_NUM_LABELS[finetuning_task]
          _SCREAMING_SNAKE_CASE							   =XLNetForSequenceClassification(_UpperCamelCase    )
     elif "squad" in finetuning_task:
          _SCREAMING_SNAKE_CASE							   =finetuning_task
          _SCREAMING_SNAKE_CASE							   =XLNetForQuestionAnswering(_UpperCamelCase    )
     else:
          _SCREAMING_SNAKE_CASE							   =XLNetLMHeadModel(_UpperCamelCase    )

     # Load weights from tf checkpoint
     load_tf_weights_in_xlnet(_UpperCamelCase							,			_UpperCamelCase							,			_UpperCamelCase    )

     # Save pytorch-model
     _SCREAMING_SNAKE_CASE							   =os.path.join(_UpperCamelCase							,			_UpperCamelCase    )
     _SCREAMING_SNAKE_CASE							   =os.path.join(_UpperCamelCase							,			_UpperCamelCase    )
     print(f"Save PyTorch model to {os.path.abspath(_UpperCamelCase    )}"    )
     torch.save(model.state_dict()							,			_UpperCamelCase    )
     print(f"Save configuration file to {os.path.abspath(_UpperCamelCase    )}"    )
     with open(_UpperCamelCase							,			'w'							,			encoding='utf-8'    ) as f:
          f.write(config.to_json_string()    )


if __name__ == "__main__":
 lowerCamelCase			:				int 			=					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(
     "--xlnet_config_file",
     default=None,
     type=str,
     required=True,
     help=(
         "The config json file corresponding to the pre-trained XLNet model. \n"
         "This specifies the model architecture."
     ),
 )
 parser.add_argument(
     "--pytorch_dump_folder_path",
     default=None,
     type=str,
     required=True,
     help="Path to the folder to store the PyTorch model or dataset/vocab.",
 )
 parser.add_argument(
     "--finetuning_task",
     default=None,
     type=str,
     help="Name of a task on which the XLNet TensorFlow model was fine-tuned",
 )
 lowerCamelCase			:				List[str] 			=					parser.parse_args()
 print(args)

 convert_xlnet_checkpoint_to_pytorch(
     args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
 )



'''simple docstring'''
import warnings


warnings.warn(
    "memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: "
    "`from accelerate import find_executable_batch_size` to avoid this warning.",
    FutureWarning,
)
