



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


UpperCAmelCase   :							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    (_lowerCAmelCase   ,	_lowerCAmelCase   ,	_lowerCAmelCase   ,	_lowerCAmelCase=None):
					# Initialise PyTorch model
					UpperCamelCase_		  =   XLNetConfig.from_json_file(_lowerCAmelCase)

					UpperCamelCase_		  =   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_		  =   finetuning_task
										UpperCamelCase_		  =   GLUE_TASKS_NUM_LABELS[finetuning_task]
										UpperCamelCase_		  =   XLNetForSequenceClassification(_lowerCAmelCase)
					elif "squad" in finetuning_task:
										UpperCamelCase_		  =   finetuning_task
										UpperCamelCase_		  =   XLNetForQuestionAnswering(_lowerCAmelCase)
					else:
										UpperCamelCase_		  =   XLNetLMHeadModel(_lowerCAmelCase)

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

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


if __name__ == "__main__":
	UpperCAmelCase   :							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""",
	)
	UpperCAmelCase   :							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
	)





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


UpperCAmelCase   :							str							 =logging.get_logger(__name__)

UpperCAmelCase   :							Tuple							 ={
    """google/pegasus-large""": """https://huggingface.co/google/pegasus-large/resolve/main/config.json""",
    # See all PEGASUS models at https://huggingface.co/models?filter=pegasus
}


class       _lowercase   (a_      ):
				'''simple docstring'''







				lowercase__   						=						"""pegasus"""
				lowercase__   						=						["""past_key_values"""]
				lowercase__   						=						{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}

				def __init__( self     , snake_case__=5_0265     , snake_case__=1024     , snake_case__=12     , snake_case__=4096     , snake_case__=16     , snake_case__=12     , snake_case__=4096     , snake_case__=16     , snake_case__=0.0     , snake_case__=0.0     , snake_case__=True     , snake_case__=True     , snake_case__="gelu"     , snake_case__=1024     , snake_case__=0.1     , snake_case__=0.0     , snake_case__=0.0     , snake_case__=0.02     , snake_case__=0     , snake_case__=False     , snake_case__=0     , snake_case__=1     , snake_case__=1     , **snake_case__     , ):

									'''simple docstring'''


									UpperCamelCase_		  =   vocab_size
									UpperCamelCase_		  =   max_position_embeddings
									UpperCamelCase_		  =   d_model
									UpperCamelCase_		  =   encoder_ffn_dim
									UpperCamelCase_		  =   encoder_layers
									UpperCamelCase_		  =   encoder_attention_heads
									UpperCamelCase_		  =   decoder_ffn_dim
									UpperCamelCase_		  =   decoder_layers
									UpperCamelCase_		  =   decoder_attention_heads
									UpperCamelCase_		  =   dropout
									UpperCamelCase_		  =   attention_dropout
									UpperCamelCase_		  =   activation_dropout
									UpperCamelCase_		  =   activation_function
									UpperCamelCase_		  =   init_std
									UpperCamelCase_		  =   encoder_layerdrop
									UpperCamelCase_		  =   decoder_layerdrop
									UpperCamelCase_		  =   use_cache
									UpperCamelCase_		  =   encoder_layers
									UpperCamelCase_		  =   scale_embedding  # scale factor will be sqrt(d_model) if True
									super().__init__(
									    pad_token_id=snake_case__     , eos_token_id=snake_case__     , is_encoder_decoder=snake_case__     , decoder_start_token_id=snake_case__     , forced_eos_token_id=snake_case__     , **snake_case__     , )

				@property
				def       _lowerCamelCase					( self    ):

									'''simple docstring'''


									return self.encoder_attention_heads

				@property
				def       _lowerCamelCase					( self    ):

									'''simple docstring'''


									return self.d_model
