






import argparse

import tensorflow as tf
import torch

from transformers import BertConfig, BertForMaskedLM
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertPooler,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


logging.set_verbosity_info()





def 						__UpperCAmelCase			(     a_							,     a_							,     a_):
							def get_masked_lm_array(a_):
														snake_case_				=							f'''masked_lm/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
														snake_case_				=							tf.train.load_variable(a_							,     a_)

														if "kernel" in name:
																					snake_case_				=							array.transpose()

														return torch.from_numpy(a_)

							def get_encoder_array(a_):
														snake_case_				=							f'''encoder/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
														snake_case_				=							tf.train.load_variable(a_							,     a_)

														if "kernel" in name:
																					snake_case_				=							array.transpose()

														return torch.from_numpy(a_)

							def get_encoder_layer_array(a_							,     a_):
														snake_case_				=							f'''encoder/_transformer_layers/{layer_index}/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
														snake_case_				=							tf.train.load_variable(a_							,     a_)

														if "kernel" in name:
																					snake_case_				=							array.transpose()

														return torch.from_numpy(a_)

							def get_encoder_attention_layer_array(a_							,     a_							,     a_):
														snake_case_				=							f'''encoder/_transformer_layers/{layer_index}/_attention_layer/{name}/.ATTRIBUTES/VARIABLE_VALUE'''
														snake_case_				=							tf.train.load_variable(a_							,     a_)
														snake_case_				=							array.reshape(a_)

														if "kernel" in name:
																					snake_case_				=							array.transpose()

														return torch.from_numpy(a_)

							print(f'''Loading model based on config from {config_path}...''')
							snake_case_				=							BertConfig.from_json_file(a_)
							snake_case_				=							BertForMaskedLM(a_)

							# Layers
							for layer_index in range(0							,     config.num_hidden_layers):
														snake_case_				=							model.bert.encoder.layer[layer_index]

														# Self-attention
														snake_case_				=							layer.attention.self

														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_query_dense/kernel'							,     self_attn.query.weight.data.shape)
														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_query_dense/bias'							,     self_attn.query.bias.data.shape)
														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_key_dense/kernel'							,     self_attn.key.weight.data.shape)
														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_key_dense/bias'							,     self_attn.key.bias.data.shape)
														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_value_dense/kernel'							,     self_attn.value.weight.data.shape)
														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_value_dense/bias'							,     self_attn.value.bias.data.shape)

														# Self-attention Output
														snake_case_				=							layer.attention.output

														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_output_dense/kernel'							,     self_output.dense.weight.data.shape)
														snake_case_				=							get_encoder_attention_layer_array(
														    a_							,     '_output_dense/bias'							,     self_output.dense.bias.data.shape)

														snake_case_				=							get_encoder_layer_array(a_							,     '_attention_layer_norm/gamma')
														snake_case_				=							get_encoder_layer_array(a_							,     '_attention_layer_norm/beta')

														# Intermediate
														snake_case_				=							layer.intermediate

														snake_case_				=							get_encoder_layer_array(a_							,     '_intermediate_dense/kernel')
														snake_case_				=							get_encoder_layer_array(a_							,     '_intermediate_dense/bias')

														# Output
														snake_case_				=							layer.output

														snake_case_				=							get_encoder_layer_array(a_							,     '_output_dense/kernel')
														snake_case_				=							get_encoder_layer_array(a_							,     '_output_dense/bias')

														snake_case_				=							get_encoder_layer_array(a_							,     '_output_layer_norm/gamma')
														snake_case_				=							get_encoder_layer_array(a_							,     '_output_layer_norm/beta')

							# Embeddings
							snake_case_				=							get_encoder_array('_position_embedding_layer/embeddings')
							snake_case_				=							get_encoder_array('_type_embedding_layer/embeddings')
							snake_case_				=							get_encoder_array('_embedding_norm_layer/gamma')
							snake_case_				=							get_encoder_array('_embedding_norm_layer/beta')

							# LM Head
							snake_case_				=							model.cls.predictions.transform

							snake_case_				=							get_masked_lm_array('dense/kernel')
							snake_case_				=							get_masked_lm_array('dense/bias')

							snake_case_				=							get_masked_lm_array('layer_norm/gamma')
							snake_case_				=							get_masked_lm_array('layer_norm/beta')

							snake_case_				=							get_masked_lm_array('embedding_table')

							# Pooling
							snake_case_				=							BertPooler(config=a_)
							snake_case_				=							get_encoder_array('_pooler_layer/kernel')
							snake_case_				=							get_encoder_array('_pooler_layer/bias')

							# Export final model
							model.save_pretrained(a_)

							# Integration test - should load without any errors ;)
							snake_case_				=							BertForMaskedLM.from_pretrained(a_)
							print(new_model.eval())

							print('Model conversion was done sucessfully!')


if __name__ == "__main__":
	lowercase  							=					argparse.ArgumentParser()
	parser.add_argument(
	    "--tf_checkpoint_path", type=str, required=True, help="Path to the TensorFlow Token Dropping checkpoint path."
	)
	parser.add_argument(
	    "--bert_config_file",
	    type=str,
	    required=True,
	    help="The config json file corresponding to the BERT model. This specifies the model architecture.",
	)
	parser.add_argument(
	    "--pytorch_dump_path",
	    type=str,
	    required=True,
	    help="Path to the output PyTorch model.",
	)
	lowercase  							=					parser.parse_args()
	convert_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)













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


lowercase  							=					{
    "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 						__UpperCAmelCase			(     a_							,     a_							,     a_							,     a_=None):
							# Initialise PyTorch model
							snake_case_				=							XLNetConfig.from_json_file(a_)

							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}''')
														snake_case_				=							finetuning_task
														snake_case_				=							GLUE_TASKS_NUM_LABELS[finetuning_task]
														snake_case_				=							XLNetForSequenceClassification(a_)
							elif "squad" in finetuning_task:
														snake_case_				=							finetuning_task
														snake_case_				=							XLNetForQuestionAnswering(a_)
							else:
														snake_case_				=							XLNetLMHeadModel(a_)

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

							# Save pytorch-model
							snake_case_				=							os.path.join(a_							,     a_)
							snake_case_				=							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__":
	lowercase  							=					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",
	)
	lowercase  							=					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
	)





