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


__UpperCamelCase	: Dict    		=  logging.get_logger(__name__)

__UpperCamelCase	: Union[str, Any]    		=  {
    '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       lowercase__		(     UpperCamelCase_):


						def __init__( self			:       Union[str, Any]					,	UpperCamelCase__			:       Union[str, Any]=None					,	UpperCamelCase__			:       Optional[int]=None					,	*UpperCamelCase__			:       Dict					,	**UpperCamelCase__			:       str		):





												'''simple docstring'''





												super().__init__(*UpperCamelCase__					,	**UpperCamelCase__		)

												if config is None:
																		assert isinstance(self.model					,	UpperCamelCase__		), (
																		    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
																		    f""" {self.model.__class__}"""
																		)
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							self.model.config
												else:
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							config

												SCREAMING_SNAKE_CASE							:					List[Any]           =							data_args
												SCREAMING_SNAKE_CASE							:					str           =							self.config.tgt_vocab_size if isinstance(self.config					,	UpperCamelCase__		) 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:
																		SCREAMING_SNAKE_CASE							:					int           =							torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id		)
												else:
																		# dynamically import label_smoothed_nll_loss
																		from utils import label_smoothed_nll_loss

																		SCREAMING_SNAKE_CASE							:					List[str]           =							label_smoothed_nll_loss


						def  __A ( self			:       Any					,	UpperCamelCase__			:       int		):





												'''simple docstring'''





												if self.optimizer is None:
																		SCREAMING_SNAKE_CASE							:					str           =							['''bias''', '''LayerNorm.weight''']
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							[
																		    {
																		        '''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,
																		    },
																		]
																		SCREAMING_SNAKE_CASE							:					int           =							Adafactor if self.args.adafactor else AdamW
																		if self.args.adafactor:
																								SCREAMING_SNAKE_CASE							:					str           =							Adafactor
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							{'''scale_parameter''': False, '''relative_step''': False}
																		else:
																								SCREAMING_SNAKE_CASE							:					Optional[int]           =							AdamW
																								SCREAMING_SNAKE_CASE							:					Any           =							{
																								    '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
																								    '''eps''': self.args.adam_epsilon,
																								}
																		SCREAMING_SNAKE_CASE							:					Dict           =							self.args.learning_rate
																		if self.sharded_ddp:
																								SCREAMING_SNAKE_CASE							:					Any           =							OSS(
																								    params=UpperCamelCase__					,	optim=UpperCamelCase__					,	**UpperCamelCase__					,	)
																		else:
																								SCREAMING_SNAKE_CASE							:					Any           =							optimizer_cls(UpperCamelCase__					,	**UpperCamelCase__		)

												if self.lr_scheduler is None:
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							self._get_lr_scheduler(UpperCamelCase__		)
												else:  # ignoring --lr_scheduler
																		logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''		)


						def  __A ( self			:       Union[str, Any]					,	UpperCamelCase__			:       List[Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							arg_to_scheduler[self.args.lr_scheduler]
												if self.args.lr_scheduler == "constant":
																		SCREAMING_SNAKE_CASE							:					List[str]           =							schedule_func(self.optimizer		)
												elif self.args.lr_scheduler == "constant_w_warmup":
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							schedule_func(self.optimizer					,	num_warmup_steps=self.args.warmup_steps		)
												else:
																		SCREAMING_SNAKE_CASE							:					str           =							schedule_func(
																		    self.optimizer					,	num_warmup_steps=self.args.warmup_steps					,	num_training_steps=UpperCamelCase__		)
												return scheduler


						def  __A ( self			:       Any		):





												'''simple docstring'''





												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  __A ( self			:       Union[str, Any]					,	UpperCamelCase__			:       Tuple					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       List[str]		):





												'''simple docstring'''





												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
																								SCREAMING_SNAKE_CASE							:					List[str]           =							model(**UpperCamelCase__					,	use_cache=UpperCamelCase__		)[0]
																								SCREAMING_SNAKE_CASE							:					List[str]           =							self.loss_fn(logits.view(-1					,	logits.shape[-1]		)					,	labels.view(-1		)		)
																		else:
																								# compute usual loss via models
																								SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							model(**UpperCamelCase__					,	labels=UpperCamelCase__					,	use_cache=UpperCamelCase__		)[:2]
												else:
																		# compute label smoothed loss
																		SCREAMING_SNAKE_CASE							:					List[str]           =							model(**UpperCamelCase__					,	use_cache=UpperCamelCase__		)[0]
																		SCREAMING_SNAKE_CASE							:					int           =							torch.nn.functional.log_softmax(UpperCamelCase__					,	dim=-1		)
																		SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE							:					int           =							self.loss_fn(UpperCamelCase__					,	UpperCamelCase__					,	self.args.label_smoothing					,	ignore_index=self.config.pad_token_id		)
												return loss, logits


						def  __A ( self			:       str					,	UpperCamelCase__			:       Dict					,	UpperCamelCase__			:       Union[str, Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							inputs.pop('''labels'''		)
												SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE							:					List[Any]           =							self._compute_loss(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)
												return loss


						def  __A ( self			:       str					,	UpperCamelCase__			:       nn.Module					,	UpperCamelCase__			:       Dict[str, Union[torch.Tensor, Any]]					,	UpperCamelCase__			:       bool					,	UpperCamelCase__			:       Optional[List[str]] = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							self._prepare_inputs(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Dict           =							{
												    '''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:
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self.model.generate(
																		    inputs['''input_ids''']					,	attention_mask=inputs['''attention_mask''']					,	**UpperCamelCase__					,	)
																		# in case the batch is shorter than max length, the output should be padded
																		if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self._pad_tensors_to_max_len(UpperCamelCase__					,	gen_kwargs['''max_length''']		)

												SCREAMING_SNAKE_CASE							:					Tuple           =							inputs.pop('''labels'''		)
												with torch.no_grad():
																		# compute loss on predict data
																		SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE							:					int           =							self._compute_loss(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							loss.mean().detach()
												if self.args.prediction_loss_only:
																		return (loss, None, None)

												SCREAMING_SNAKE_CASE							:					int           =							generated_tokens if self.args.predict_with_generate else logits

												if labels.shape[-1] < gen_kwargs["max_length"]:
																		SCREAMING_SNAKE_CASE							:					Dict           =							self._pad_tensors_to_max_len(UpperCamelCase__					,	gen_kwargs['''max_length''']		)

												return (loss, logits, labels)



						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							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}"""		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							pad_token_id * torch.ones(
												    (tensor.shape[0], max_length)					,	dtype=tensor.dtype					,	device=tensor.device		)
												SCREAMING_SNAKE_CASE							:					Tuple           =							tensor
												return padded_tensor






from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


if TYPE_CHECKING:
					from ... import FeatureExtractionMixin, PreTrainedTokenizerBase, TensorType


__UpperCamelCase	: Optional[Any]    		=  logging.get_logger(__name__)

__UpperCamelCase	: Optional[Any]    		=  {
    'microsoft/deberta-v2-xlarge': 'https://huggingface.co/microsoft/deberta-v2-xlarge/resolve/main/config.json',
    'microsoft/deberta-v2-xxlarge': 'https://huggingface.co/microsoft/deberta-v2-xxlarge/resolve/main/config.json',
    'microsoft/deberta-v2-xlarge-mnli': (
        'https://huggingface.co/microsoft/deberta-v2-xlarge-mnli/resolve/main/config.json'
    ),
    'microsoft/deberta-v2-xxlarge-mnli': (
        'https://huggingface.co/microsoft/deberta-v2-xxlarge-mnli/resolve/main/config.json'
    ),
}





class       lowercase__		(     UpperCamelCase_):
						UpperCamelCase_  =	"""deberta-v2"""


						def __init__( self			:       Optional[Any]					,	UpperCamelCase__			:       Any=12_8100					,	UpperCamelCase__			:       Optional[int]=1536					,	UpperCamelCase__			:       Dict=24					,	UpperCamelCase__			:       List[str]=24					,	UpperCamelCase__			:       Tuple=6144					,	UpperCamelCase__			:       List[Any]="gelu"					,	UpperCamelCase__			:       Optional[int]=0.1					,	UpperCamelCase__			:       List[str]=0.1					,	UpperCamelCase__			:       Optional[int]=512					,	UpperCamelCase__			:       Optional[Any]=0					,	UpperCamelCase__			:       Union[str, Any]=0.02					,	UpperCamelCase__			:       List[Any]=1E-7					,	UpperCamelCase__			:       List[str]=False					,	UpperCamelCase__			:       str=-1					,	UpperCamelCase__			:       List[Any]=0					,	UpperCamelCase__			:       List[Any]=True					,	UpperCamelCase__			:       Optional[Any]=None					,	UpperCamelCase__			:       Tuple=0					,	UpperCamelCase__			:       str="gelu"					,	**UpperCamelCase__			:       Optional[int]					,	):





												'''simple docstring'''





												super().__init__(**UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							hidden_size
												SCREAMING_SNAKE_CASE							:					int           =							num_hidden_layers
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							num_attention_heads
												SCREAMING_SNAKE_CASE							:					List[Any]           =							intermediate_size
												SCREAMING_SNAKE_CASE							:					Any           =							hidden_act
												SCREAMING_SNAKE_CASE							:					Dict           =							hidden_dropout_prob
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							attention_probs_dropout_prob
												SCREAMING_SNAKE_CASE							:					str           =							max_position_embeddings
												SCREAMING_SNAKE_CASE							:					Dict           =							type_vocab_size
												SCREAMING_SNAKE_CASE							:					str           =							initializer_range
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							relative_attention
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							max_relative_positions
												SCREAMING_SNAKE_CASE							:					Optional[int]           =							pad_token_id
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							position_biased_input

												# Backwards compatibility
												if type(UpperCamelCase__		) == str:
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							[x.strip() for x in pos_att_type.lower().split('''|'''		)]

												SCREAMING_SNAKE_CASE							:					Any           =							pos_att_type
												SCREAMING_SNAKE_CASE							:					int           =							vocab_size
												SCREAMING_SNAKE_CASE							:					Dict           =							layer_norm_eps

												SCREAMING_SNAKE_CASE							:					Tuple           =							kwargs.get('''pooler_hidden_size'''					,	UpperCamelCase__		)
												SCREAMING_SNAKE_CASE							:					Any           =							pooler_dropout
												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							pooler_hidden_act






class       lowercase__		(     UpperCamelCase_):


						@property
						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												if self.task == "multiple-choice":
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
												else:
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							{0: '''batch''', 1: '''sequence'''}
												if self._config.type_vocab_size > 0:
																		return OrderedDict(
																		    [('''input_ids''', dynamic_axis), ('''attention_mask''', dynamic_axis), ('''token_type_ids''', dynamic_axis)]		)
												else:
																		return OrderedDict([('''input_ids''', dynamic_axis), ('''attention_mask''', dynamic_axis)]		)


						@property
						def  __A ( self			:       Union[str, Any]		):





												'''simple docstring'''





												return 12



						def  __A ( self			:       Dict					,	UpperCamelCase__			:       Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]					,	UpperCamelCase__			:       int = -1					,	UpperCamelCase__			:       int = -1					,	UpperCamelCase__			:       int = -1					,	UpperCamelCase__			:       bool = False					,	UpperCamelCase__			:       Optional["TensorType"] = None					,	UpperCamelCase__			:       int = 3					,	UpperCamelCase__			:       int = 40					,	UpperCamelCase__			:       int = 40					,	UpperCamelCase__			:       "PreTrainedTokenizerBase" = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							super().generate_dummy_inputs(preprocessor=UpperCamelCase__					,	framework=UpperCamelCase__		)
												if self._config.type_vocab_size == 0 and "token_type_ids" in dummy_inputs:
																		del dummy_inputs["token_type_ids"]
												return dummy_inputs





