




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


lowerCAmelCase_					    = logging.get_logger(__name__)

lowerCAmelCase_					    = {
    """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  _lowerCAmelCase     (  __UpperCAmelCase	):

      def __init__(   self      ,				__UpperCAmelCase=None      ,				__UpperCAmelCase=None      ,				*__UpperCAmelCase      ,				**__UpperCAmelCase		):
        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__}"""
          )
          lowerCAmelCase__				:  Dict		       =     self.model.config
        else:
          lowerCAmelCase__				:  Optional[Any]		       =     config

        lowerCAmelCase__				:  Any		       =     data_args
        lowerCAmelCase__				:  Dict		       =     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:
          lowerCAmelCase__				:  str		       =     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__				:  Optional[Any]		       =     label_smoothed_nll_loss

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):

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

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

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  List[Any]		       =     arg_to_scheduler[self.args.lr_scheduler]
        if self.args.lr_scheduler == "constant":
          lowerCAmelCase__				:  int		       =     schedule_func(self.optimizer		)
        elif self.args.lr_scheduler == "constant_w_warmup":
          lowerCAmelCase__				:  List[str]		       =     schedule_func(self.optimizer      ,				num_warmup_steps=self.args.warmup_steps		)
        else:
          lowerCAmelCase__				:  Tuple		       =     schedule_func(
              self.optimizer      ,				num_warmup_steps=self.args.warmup_steps      ,				num_training_steps=__UpperCAmelCase		)
        return scheduler

      def 	__magic_name__(   self		):

        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 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):

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

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  int		       =     inputs.pop('''labels'''		)
        lowerCAmelCase__       ,			lowerCAmelCase__				:  Dict		       =     self._compute_loss(__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		)
        return loss

      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase = None      ,				):
        lowerCAmelCase__				:  int		       =     self._prepare_inputs(__UpperCAmelCase		)

        lowerCAmelCase__				:  Union[str, Any]		       =     {
            '''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__				:  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"]:
            lowerCAmelCase__				:  Union[str, Any]		       =     self._pad_tensors_to_max_len(__UpperCAmelCase      ,				gen_kwargs['''max_length''']		)

        lowerCAmelCase__				:  Any		       =     inputs.pop('''labels'''		)
        with torch.no_grad():
          # compute loss on predict data
          lowerCAmelCase__       ,			lowerCAmelCase__				:  List[str]		       =     self._compute_loss(__UpperCAmelCase      ,				__UpperCAmelCase      ,				__UpperCAmelCase		)

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

        lowerCAmelCase__				:  str		       =     generated_tokens if self.args.predict_with_generate else logits

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

        return (loss, logits, labels)




      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  List[str]		       =     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__				:  int		       =     pad_token_id * torch.ones(
            (tensor.shape[0], max_length)      ,				dtype=tensor.dtype      ,				device=tensor.device		)
        lowerCAmelCase__				:  Optional[int]		       =     tensor
        return padded_tensor







import warnings

from ...utils import logging
from .image_processing_owlvit import OwlViTImageProcessor


lowerCAmelCase_					    = logging.get_logger(__name__)






class  _lowerCAmelCase     (  _lowercase	):

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

