






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 				_A			(     UpperCamelCase				):
    """simple docstring"""

    def __init__(				self      :     List[str] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any]=None ,     __SCREAMING_SNAKE_CASE      :     str=None ,     *__SCREAMING_SNAKE_CASE      :     Union[str, Any] ,     **__SCREAMING_SNAKE_CASE      :     List[Any]      )      ->							Any:
           super().__init__(*__SCREAMING_SNAKE_CASE ,     **__SCREAMING_SNAKE_CASE      )

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

           __UpperCAmelCase     				=data_args
           __UpperCAmelCase     				=self.config.tgt_vocab_size if isinstance(self.config ,     __SCREAMING_SNAKE_CASE      ) 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:
                  __UpperCAmelCase     				=torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id      )
           else:
                  # dynamically import label_smoothed_nll_loss
                  from utils import label_smoothed_nll_loss

                  __UpperCAmelCase     				=label_smoothed_nll_loss



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

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



    def _a							(				self      :     Optional[Any] ,     __SCREAMING_SNAKE_CASE      :     Optional[int]      )      ->							Any:
           __UpperCAmelCase     				=arg_to_scheduler[self.args.lr_scheduler]
           if self.args.lr_scheduler == "constant":
                  __UpperCAmelCase     				=schedule_func(self.optimizer      )
           elif self.args.lr_scheduler == "constant_w_warmup":
                  __UpperCAmelCase     				=schedule_func(self.optimizer ,     num_warmup_steps=self.args.warmup_steps      )
           else:
                  __UpperCAmelCase     				=schedule_func(
                      self.optimizer ,     num_warmup_steps=self.args.warmup_steps ,     num_training_steps=__SCREAMING_SNAKE_CASE      )
           return scheduler



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



    def _a							(				self      :     Optional[Any] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any] ,     __SCREAMING_SNAKE_CASE      :     List[Any]      )      ->							Dict:
           __UpperCAmelCase     				=inputs.pop("""labels"""      )
           __UpperCAmelCase				, __UpperCAmelCase     				=self._compute_loss(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      )
           return loss



    def _a							(				self      :     List[str] ,     __SCREAMING_SNAKE_CASE      :     nn.Module ,     __SCREAMING_SNAKE_CASE      :     Dict[str, Union[torch.Tensor, Any]] ,     __SCREAMING_SNAKE_CASE      :     bool ,     __SCREAMING_SNAKE_CASE      :     Optional[List[str]] = None ,     )      ->							Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
           __UpperCAmelCase     				=self._prepare_inputs(__SCREAMING_SNAKE_CASE      )

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

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

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

           __UpperCAmelCase     				=generated_tokens if self.args.predict_with_generate else logits

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

           return (loss, logits, labels)



    def _a							(				self      :     Optional[Any] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any] ,     __SCREAMING_SNAKE_CASE      :     int      )      ->							List[Any]:
           # If PAD token is not defined at least EOS token has to be defined
           __UpperCAmelCase     				=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}'''      )

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









'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


snake_case_					:      Tuple    			=    {
    'configuration_layoutlmv2': ['LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP', 'LayoutLMv2Config'],
    'processing_layoutlmv2': ['LayoutLMv2Processor'],
    'tokenization_layoutlmv2': ['LayoutLMv2Tokenizer'],
}

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_					:      int    			=    ['LayoutLMv2TokenizerFast']

try:
					if not is_vision_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_					:      Union[str, Any]    			=    ['LayoutLMv2FeatureExtractor']
					snake_case_					:      str    			=    ['LayoutLMv2ImageProcessor']

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_					:      Optional[int]    			=    [
					    'LAYOUTLMV2_PRETRAINED_MODEL_ARCHIVE_LIST',
					    'LayoutLMv2ForQuestionAnswering',
					    'LayoutLMv2ForSequenceClassification',
					    'LayoutLMv2ForTokenClassification',
					    'LayoutLMv2Layer',
					    'LayoutLMv2Model',
					    'LayoutLMv2PreTrainedModel',
					]

if TYPE_CHECKING:
					from .configuration_layoutlmva import LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMvaConfig
					from .processing_layoutlmva import LayoutLMvaProcessor
					from .tokenization_layoutlmva import LayoutLMvaTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_layoutlmva_fast import LayoutLMvaTokenizerFast

					try:
										if not is_vision_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .feature_extraction_layoutlmva import LayoutLMvaFeatureExtractor, LayoutLMvaImageProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_layoutlmva import (
										    LAYOUTLMV2_PRETRAINED_MODEL_ARCHIVE_LIST,
										    LayoutLMvaForQuestionAnswering,
										    LayoutLMvaForSequenceClassification,
										    LayoutLMvaForTokenClassification,
										    LayoutLMvaLayer,
										    LayoutLMvaModel,
										    LayoutLMvaPreTrainedModel,
										)
else:
					import sys

					snake_case_					:      int    			=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


