





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  :    Optional[int]						    =  logging.get_logger(__name__)

_snake_case  :    str						    =  {
    '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    (    __snake_case    ):

  """simple docstring"""


  def __init__(				self :		Union[str, Any]					,  lowerCAmelCase_ :		Optional[Any]=None					,  lowerCAmelCase_ :		str=None					,  *lowerCAmelCase_ :		List[str]					,  **lowerCAmelCase_ :		Union[str, Any]     )							->						Optional[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 :		List[Any]					,  lowerCAmelCase_ :		Union[str, Any]     )							->						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=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 :		str					,  lowerCAmelCase_ :		Dict     )							->						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=lowerCAmelCase_     )
     return scheduler






  def 				lowercase		(				self :		List[Any]     )							->						Optional[Any]:
     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 :		int					,  lowerCAmelCase_ :		str					,  lowerCAmelCase_ :		Any					,  lowerCAmelCase_ :		Optional[int]     )							->						Optional[int]:
     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								=	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								=	self.loss_fn(lowerCAmelCase_					,  lowerCAmelCase_					,  self.args.label_smoothing					,  ignore_index=self.config.pad_token_id     )
     return loss, logits






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






  def 				lowercase		(				self :		int					,  lowerCAmelCase_ :		Dict					,  lowerCAmelCase_ :		Tuple					,  lowerCAmelCase_ :		Optional[int]					,  lowerCAmelCase_ :		Tuple = None					,  )							->						int:
     __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								=	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 :		Optional[int]					,  lowerCAmelCase_ :		Optional[Any]					,  lowerCAmelCase_ :		str     )							->						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 warnings

from ...utils import logging
from .image_processing_dpt import DPTImageProcessor


snake_case_			  = logging.get_logger(__name__)


class       SCREAMING_SNAKE_CASE__					(__snake_case     ):





    def __init__(      self						,   *a						,   **a):
         warnings.warn(
             'The class DPTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please'
             ' use DPTImageProcessor instead.'						,   a						,   )
         super().__init__(*a						,   **a)

