



"""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[Any]	=   {
    '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__						:						Tuple	,			UpperCamelCase__						:						Union[str, Any]	,			UpperCamelCase__						:						str	,			UpperCamelCase__						:						str=None				):
    # Initialise PyTorch model
    _UpperCAmelCase				:    Optional[int]											=       XLNetConfig.from_json_file(A__				)

    _UpperCAmelCase				:    Optional[int]											=       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}'				)
        _UpperCAmelCase				:    Dict											=       finetuning_task
        _UpperCAmelCase				:    Union[str, Any]											=       GLUE_TASKS_NUM_LABELS[finetuning_task]
        _UpperCAmelCase				:    Optional[int]											=       XLNetForSequenceClassification(A__				)
    elif "squad" in finetuning_task:
        _UpperCAmelCase				:    Optional[int]											=       finetuning_task
        _UpperCAmelCase				:    List[Any]											=       XLNetForQuestionAnswering(A__				)
    else:
        _UpperCAmelCase				:    Union[str, Any]											=       XLNetLMHeadModel(A__				)

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

    # Save pytorch-model
    _UpperCAmelCase				:    Dict											=       os.path.join(A__	,			A__				)
    _UpperCAmelCase				:    Union[str, Any]											=       os.path.join(A__	,			A__				)
    print(F'Save PyTorch model to {os.path.abspath(A__				)}'				)
    torch.save(model.state_dict()	,			A__				)
    print(F'Save configuration file to {os.path.abspath(A__				)}'				)
    with open(A__	,			'''w'''	,			encoding='''utf-8'''				) as f:
        f.write(config.to_json_string()				)


if __name__ == "__main__":
  _lowerCAmelCase       :Optional[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[Any]	=   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 math






def      UpperCamelCase_				(      A__						):
	return math.sqrt(A__						) * math.sqrt(A__						) == num






def      UpperCamelCase_				(      A__						):
	a_      				=						0
	a_      				=						n
	while left <= right:
		a_      				=						(left + right) // 2
		if mid**2 == n:
			return True
		elif mid**2 > n:
			a_      				=						mid - 1
		else:
			a_      				=						mid + 1
	return False


if __name__ == "__main__":
		import doctest

		doctest.testmod()
