





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


_snake_case  :   Union[str, Any]						   =  logging.get_logger(__name__)

_snake_case  :   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     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  def __init__(    self			:		List[str] ,  lowerCAmelCase_			:		List[Any]=None ,  lowerCAmelCase_			:		Dict=None ,  *lowerCAmelCase_			:		Tuple ,  **lowerCAmelCase_			:		Any					)							->						Any:
     super().__init__(*lowerCAmelCase_ ,  **lowerCAmelCase_					)

     if config is None:
        assert isinstance(self.model ,  lowerCAmelCase_					), (
            "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 ,  lowerCAmelCase_					) 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     lowercase		(    self			:		Tuple ,  lowerCAmelCase_			:		int					)							->						int:
     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=lowerCAmelCase_ ,  optim=lowerCAmelCase_ ,  **lowerCAmelCase_ ,  )
        else:
           __lowerCAmelCase							     =     optimizer_cls(lowerCAmelCase_ ,  **lowerCAmelCase_					)

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






  def     lowercase		(    self			:		Tuple ,  lowerCAmelCase_			:		Union[str, Any]					)							->						List[Any]:
     __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=lowerCAmelCase_					)
     return scheduler






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






  def     lowercase		(    self			:		List[Any] ,  lowerCAmelCase_			:		List[str] ,  lowerCAmelCase_			:		Optional[Any]					)							->						Optional[Any]:
     __lowerCAmelCase							     =     inputs.pop('labels'					)
     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self._compute_loss(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)
     return loss






  def     lowercase		(    self			:		int ,  lowerCAmelCase_			:		nn.Module ,  lowerCAmelCase_			:		Dict[str, Union[torch.Tensor, Any]] ,  lowerCAmelCase_			:		bool ,  lowerCAmelCase_			:		Optional[List[str]] = None ,  )							->						Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
     __lowerCAmelCase							     =     self._prepare_inputs(lowerCAmelCase_					)

     __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'] ,  **lowerCAmelCase_ ,  )
        # 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(lowerCAmelCase_ ,  gen_kwargs['max_length']					)

     __lowerCAmelCase							     =     inputs.pop('labels'					)
     with torch.no_grad():
        # compute loss on predict data
        __lowerCAmelCase     ,     __lowerCAmelCase							     =     self._compute_loss(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)

     __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(lowerCAmelCase_ ,  gen_kwargs['max_length']					)

     return (loss, logits, labels)






  def     lowercase		(    self			:		List[str] ,  lowerCAmelCase_			:		Union[str, Any] ,  lowerCAmelCase_			:		List[Any]					)							->						Optional[int]:
     # 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








from typing import TYPE_CHECKING

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


_snake_case  :   Optional[Any]						   =  {'configuration_fnet': ['FNET_PRETRAINED_CONFIG_ARCHIVE_MAP', 'FNetConfig']}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _snake_case  :   int						   =  ['FNetTokenizer']

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _snake_case  :   Tuple						   =  ['FNetTokenizerFast']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _snake_case  :   List[str]						   =  [
         'FNET_PRETRAINED_MODEL_ARCHIVE_LIST',
         'FNetForMaskedLM',
         'FNetForMultipleChoice',
         'FNetForNextSentencePrediction',
         'FNetForPreTraining',
         'FNetForQuestionAnswering',
         'FNetForSequenceClassification',
         'FNetForTokenClassification',
         'FNetLayer',
         'FNetModel',
         'FNetPreTrainedModel',
     ]


if TYPE_CHECKING:
     from .configuration_fnet import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FNetConfig

     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_fnet import FNetTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_fnet_fast import FNetTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_fnet import (
              FNET_PRETRAINED_MODEL_ARCHIVE_LIST,
              FNetForMaskedLM,
              FNetForMultipleChoice,
              FNetForNextSentencePrediction,
              FNetForPreTraining,
              FNetForQuestionAnswering,
              FNetForSequenceClassification,
              FNetForTokenClassification,
              FNetLayer,
              FNetModel,
              FNetPreTrainedModel,
          )


else:
     import sys

     _snake_case  :   str						   =  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

