
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
			from fairscale.optim import OSS


__A							=logging.get_logger(__name__)

__A							={
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    '''constant''': get_constant_schedule,
    '''constant_w_warmup''': get_constant_schedule_with_warmup,
}




class        _SCREAMING_SNAKE_CASE	( snake_case_		):



	def __init__(       self		,    lowercase=None		,    lowercase=None		,    *lowercase		,    **lowercase					)							->			Tuple:
					super().__init__(*lowercase		,    **lowercase					)

					if config is None:
									assert isinstance(self.model		,    lowercase					), (
									    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
									    f' {self.model.__class__}'
									)
									lowerCamelCase_         =     self.model.config
					else:
									lowerCamelCase_         =     config

					lowerCamelCase_         =     data_args
					lowerCamelCase_         =     self.config.tgt_vocab_size if isinstance(self.config		,    lowercase					) else self.config.vocab_size

					if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
									assert self.config.pad_token_id is not None, (
									    "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
									    " calculation or doing label smoothing."
									)

					if self.config.pad_token_id is None and self.config.eos_token_id is not None:
									logger.warning(
									    f'The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'
									    " padding.."					)

					if self.args.label_smoothing == 0:
									lowerCamelCase_         =     torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id					)
					else:
									# dynamically import label_smoothed_nll_loss
									from utils import label_smoothed_nll_loss

									lowerCamelCase_         =     label_smoothed_nll_loss



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase					)							->			Union[str, Any]:
					if self.optimizer is None:
									lowerCamelCase_         =     ["bias", "LayerNorm.weight"]
									lowerCamelCase_         =     [
									    {
									        "params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay					)],
									        "weight_decay": self.args.weight_decay,
									    },
									    {
									        "params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay					)],
									        "weight_decay": 0.0,
									    },
									]
									lowerCamelCase_         =     Adafactor if self.args.adafactor else AdamW
									if self.args.adafactor:
													lowerCamelCase_         =     Adafactor
													lowerCamelCase_         =     {"scale_parameter": False, "relative_step": False}
									else:
													lowerCamelCase_         =     AdamW
													lowerCamelCase_         =     {
													    "betas": (self.args.adam_betaa, self.args.adam_betaa),
													    "eps": self.args.adam_epsilon,
													}
									lowerCamelCase_         =     self.args.learning_rate
									if self.sharded_ddp:
													lowerCamelCase_         =     OSS(
													    params=lowercase		,    optim=lowercase		,    **lowercase		,    )
									else:
													lowerCamelCase_         =     optimizer_cls(lowercase		,    **lowercase					)

					if self.lr_scheduler is None:
									lowerCamelCase_         =     self._get_lr_scheduler(lowercase					)
					else:  # ignoring --lr_scheduler
									logger.warning("scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored."					)



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase					)							->			Tuple:
					lowerCamelCase_         =     arg_to_scheduler[self.args.lr_scheduler]
					if self.args.lr_scheduler == "constant":
									lowerCamelCase_         =     schedule_func(self.optimizer					)
					elif self.args.lr_scheduler == "constant_w_warmup":
									lowerCamelCase_         =     schedule_func(self.optimizer		,    num_warmup_steps=self.args.warmup_steps					)
					else:
									lowerCamelCase_         =     schedule_func(
									    self.optimizer		,    num_warmup_steps=self.args.warmup_steps		,    num_training_steps=lowercase					)
					return scheduler



	def        SCREAMING_SNAKE_CASE_(       self					)							->			Optional[torch.utils.data.Sampler]:
					if isinstance(self.train_dataset		,    torch.utils.data.IterableDataset					):
									return None
					elif is_torch_tpu_available():
									return get_tpu_sampler(self.train_dataset					)
					else:
									if self.args.sortish_sampler:
													self.train_dataset.make_sortish_sampler(
													    self.args.per_device_train_batch_size		,    distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)		,    )

									return (
									    RandomSampler(self.train_dataset					)
									    if self.args.local_rank == -1
									    else DistributedSampler(self.train_dataset					)
									)



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase		,    lowercase					)							->			List[str]:
					if self.args.label_smoothing == 0:
									if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
													# force training to ignore pad token
													lowerCamelCase_         =     model(**lowercase		,    use_cache=lowercase					)[0]
													lowerCamelCase_         =     self.loss_fn(logits.view(-1		,    logits.shape[-1]					)		,    labels.view(-1					)					)
									else:
													# compute usual loss via models
													lowerCamelCase_      ,       lowerCamelCase_         =     model(**lowercase		,    labels=lowercase		,    use_cache=lowercase					)[:2]
					else:
									# compute label smoothed loss
									lowerCamelCase_         =     model(**lowercase		,    use_cache=lowercase					)[0]
									lowerCamelCase_         =     torch.nn.functional.log_softmax(lowercase		,    dim=-1					)
									lowerCamelCase_      ,       lowerCamelCase_         =     self.loss_fn(lowercase		,    lowercase		,    self.args.label_smoothing		,    ignore_index=self.config.pad_token_id					)
					return loss, logits



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase					)							->			Tuple:
					lowerCamelCase_         =     inputs.pop("labels"					)
					lowerCamelCase_      ,       lowerCamelCase_         =     self._compute_loss(lowercase		,    lowercase		,    lowercase					)
					return loss



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase		,    lowercase		,    lowercase = None		,    )							->			Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
					lowerCamelCase_         =     self._prepare_inputs(lowercase					)

					lowerCamelCase_         =     {
					    "max_length": self.data_args.val_max_target_length
					    if self.data_args is not None
					    else self.config.max_length,
					    "num_beams": self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
					}

					if self.args.predict_with_generate and not self.args.prediction_loss_only:
									lowerCamelCase_         =     self.model.generate(
									    inputs["input_ids"]		,    attention_mask=inputs["attention_mask"]		,    **lowercase		,    )
									# in case the batch is shorter than max length, the output should be padded
									if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
													lowerCamelCase_         =     self._pad_tensors_to_max_len(lowercase		,    gen_kwargs["max_length"]					)

					lowerCamelCase_         =     inputs.pop("labels"					)
					with torch.no_grad():
									# compute loss on predict data
									lowerCamelCase_      ,       lowerCamelCase_         =     self._compute_loss(lowercase		,    lowercase		,    lowercase					)

					lowerCamelCase_         =     loss.mean().detach()
					if self.args.prediction_loss_only:
									return (loss, None, None)

					lowerCamelCase_         =     generated_tokens if self.args.predict_with_generate else logits

					if labels.shape[-1] < gen_kwargs["max_length"]:
									lowerCamelCase_         =     self._pad_tensors_to_max_len(lowercase		,    gen_kwargs["max_length"]					)

					return (loss, logits, labels)



	def        SCREAMING_SNAKE_CASE_(       self		,    lowercase		,    lowercase					)							->			Any:
					# If PAD token is not defined at least EOS token has to be defined
					lowerCamelCase_         =     self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

					if pad_token_id is None:
									raise ValueError(
									    "Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be"
									    f' padded to `max_length`={max_length}'					)

					lowerCamelCase_         =     pad_token_id * torch.ones(
					    (tensor.shape[0], max_length)		,    dtype=tensor.dtype		,    device=tensor.device					)
					lowerCamelCase_         =     tensor
					return padded_tensor




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


__A							={
    '''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
				lowerCamelCase_         =     XLNetConfig.from_json_file(lowerCamelCase__       )

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

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

				# Save pytorch-model
				lowerCamelCase_         =     os.path.join(lowerCamelCase__      ,					lowerCamelCase__       )
				lowerCamelCase_         =     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__":
			__A							=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''',
			)
			__A							=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
			)


