

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         =  logging.get_logger(__name__)

_snake_case         =  {
    """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  ):


   def __init__(			self      ,    _UpperCamelCase=None      ,    _UpperCamelCase=None      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """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__}'''
         )
         _lowercase      :				Dict      					=       self.model.config
      else:
         _lowercase      :				List[str]      					=       config

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

         _lowercase      :				Any      					=       label_smoothed_nll_loss


   def 						_lowerCamelCase       (			self      ,    _UpperCamelCase  ):





      """simple docstring"""





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

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


   def 						_lowerCamelCase       (			self      ,    _UpperCamelCase  ):





      """simple docstring"""





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


   def 						_lowerCamelCase       (			self  ):





      """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 						_lowerCamelCase       (			self      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase  ):





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


   def 						_lowerCamelCase       (			self      ,    _UpperCamelCase      ,    _UpperCamelCase  ):





      """simple docstring"""





      _lowercase      :				List[Any]      					=       inputs.pop("labels"  )
      _lowercase      :				Optional[Any]      					=       self._compute_loss(_UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase  )
      return loss


   def 						_lowerCamelCase       (			self      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase = None      ,    ):





      """simple docstring"""





      _lowercase      :				Any      					=       self._prepare_inputs(_UpperCamelCase  )

      _lowercase      :				str      					=       {
          '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:
         _lowercase      :				Optional[int]      					=       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"]:
            _lowercase      :				Dict      					=       self._pad_tensors_to_max_len(_UpperCamelCase      ,    gen_kwargs["max_length"]  )

      _lowercase      :				Optional[int]      					=       inputs.pop("labels"  )
      with torch.no_grad():
         # compute loss on predict data
         _lowercase      :				Any      					=       self._compute_loss(_UpperCamelCase      ,    _UpperCamelCase      ,    _UpperCamelCase  )

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

      _lowercase      :				Dict      					=       generated_tokens if self.args.predict_with_generate else logits

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

      return (loss, logits, labels)






   def 						_lowerCamelCase       (			self      ,    _UpperCamelCase      ,    _UpperCamelCase  ):





      """simple docstring"""





      _lowercase      :				Optional[int]      					=       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}'''  )

      _lowercase      :				Optional[int]      					=       pad_token_id * torch.ones(
          (tensor.shape[0], max_length)      ,    dtype=tensor.dtype      ,    device=tensor.device  )
      _lowercase      :				Optional[Any]      					=       tensor
      return padded_tensor









'''simple docstring'''
import warnings
from typing import List

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import is_flax_available, is_tf_available, is_torch_available


class        a__  (							lowerCamelCase_  ):
   _SCREAMING_SNAKE_CASE   :					List[Any]     	=							['image_processor', 'tokenizer']
   _SCREAMING_SNAKE_CASE   :					str     	=							'OwlViTImageProcessor'
   _SCREAMING_SNAKE_CASE   :					List[str]     	=							('CLIPTokenizer', 'CLIPTokenizerFast')


   def __init__(			self      ,    _UpperCamelCase=None      ,    _UpperCamelCase=None      ,    **_UpperCamelCase  ):





      """simple docstring"""





      _lowercase      :				Dict      					=       None
      if "feature_extractor" in kwargs:
         warnings.warn(
             "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
             " instead."      ,    _UpperCamelCase      ,    )
         _lowercase      :				Optional[int]      					=       kwargs.pop("feature_extractor"  )

      _lowercase      :				str      					=       image_processor if image_processor is not None else feature_extractor
      if image_processor is None:
         raise ValueError("You need to specify an `image_processor`."  )
      if tokenizer is None:
         raise ValueError("You need to specify a `tokenizer`."  )

      super().__init__(_UpperCamelCase      ,    _UpperCamelCase  )


   def __call__(			self      ,    _UpperCamelCase=None      ,    _UpperCamelCase=None      ,    _UpperCamelCase=None      ,    _UpperCamelCase="max_length"      ,    _UpperCamelCase="np"      ,    **_UpperCamelCase  ):





      """simple docstring"""





      if text is None and query_images is None and images is None:
         raise ValueError(
             "You have to specify at least one text or query image or image. All three cannot be none."  )

      if text is not None:
         if isinstance(_UpperCamelCase      ,    _UpperCamelCase  ) or (isinstance(_UpperCamelCase      ,    _UpperCamelCase  ) and not isinstance(text[0]      ,    _UpperCamelCase  )):
            _lowercase      :				int      					=       [self.tokenizer(_UpperCamelCase      ,    padding=_UpperCamelCase      ,    return_tensors=_UpperCamelCase      ,    **_UpperCamelCase  )]

         elif isinstance(_UpperCamelCase      ,    _UpperCamelCase  ) and isinstance(text[0]      ,    _UpperCamelCase  ):
            _lowercase      :				str      					=       []

            # Maximum number of queries across batch
            _lowercase      :				str      					=       max([len(_UpperCamelCase  ) for t in text]  )

            # Pad all batch samples to max number of text queries
            for t in text:
               if len(_UpperCamelCase  ) != max_num_queries:
                  _lowercase      :				List[Any]      					=       t + [" "] * (max_num_queries - len(_UpperCamelCase  ))

               _lowercase      :				Tuple      					=       self.tokenizer(_UpperCamelCase      ,    padding=_UpperCamelCase      ,    return_tensors=_UpperCamelCase      ,    **_UpperCamelCase  )
               encodings.append(_UpperCamelCase  )
         else:
            raise TypeError("Input text should be a string, a list of strings or a nested list of strings"  )

         if return_tensors == "np":
            _lowercase      :				List[Any]      					=       np.concatenate([encoding["input_ids"] for encoding in encodings]      ,    axis=0  )
            _lowercase      :				Optional[Any]      					=       np.concatenate([encoding["attention_mask"] for encoding in encodings]      ,    axis=0  )

         elif return_tensors == "jax" and is_flax_available():
            import jax.numpy as jnp

            _lowercase      :				Union[str, Any]      					=       jnp.concatenate([encoding["input_ids"] for encoding in encodings]      ,    axis=0  )
            _lowercase      :				int      					=       jnp.concatenate([encoding["attention_mask"] for encoding in encodings]      ,    axis=0  )

         elif return_tensors == "pt" and is_torch_available():
            import torch

            _lowercase      :				int      					=       torch.cat([encoding["input_ids"] for encoding in encodings]      ,    dim=0  )
            _lowercase      :				Dict      					=       torch.cat([encoding["attention_mask"] for encoding in encodings]      ,    dim=0  )

         elif return_tensors == "tf" and is_tf_available():
            import tensorflow as tf

            _lowercase      :				Optional[int]      					=       tf.stack([encoding["input_ids"] for encoding in encodings]      ,    axis=0  )
            _lowercase      :				List[str]      					=       tf.stack([encoding["attention_mask"] for encoding in encodings]      ,    axis=0  )

         else:
            raise ValueError("Target return tensor type could not be returned"  )

         _lowercase      :				Optional[int]      					=       BatchEncoding()
         _lowercase      :				List[Any]      					=       input_ids
         _lowercase      :				Dict      					=       attention_mask

      if query_images is not None:
         _lowercase      :				int      					=       BatchEncoding()
         _lowercase      :				Any      					=       self.image_processor(
             _UpperCamelCase      ,    return_tensors=_UpperCamelCase      ,    **_UpperCamelCase  ).pixel_values
         _lowercase      :				Any      					=       query_pixel_values

      if images is not None:
         _lowercase      :				str      					=       self.image_processor(_UpperCamelCase      ,    return_tensors=_UpperCamelCase      ,    **_UpperCamelCase  )

      if text is not None and images is not None:
         _lowercase      :				List[Any]      					=       image_features.pixel_values
         return encoding
      elif query_images is not None and images is not None:
         _lowercase      :				Optional[Any]      					=       image_features.pixel_values
         return encoding
      elif text is not None or query_images is not None:
         return encoding
      else:
         return BatchEncoding(data=dict(**_UpperCamelCase  )      ,    tensor_type=_UpperCamelCase  )


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      return self.image_processor.post_process(*_UpperCamelCase      ,    **_UpperCamelCase  )


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      return self.image_processor.post_process_object_detection(*_UpperCamelCase      ,    **_UpperCamelCase  )


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      return self.image_processor.post_process_image_guided_detection(*_UpperCamelCase      ,    **_UpperCamelCase  )


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      return self.tokenizer.batch_decode(*_UpperCamelCase      ,    **_UpperCamelCase  )


   def 						_lowerCamelCase       (			self      ,    *_UpperCamelCase      ,    **_UpperCamelCase  ):





      """simple docstring"""





      return self.tokenizer.decode(*_UpperCamelCase      ,    **_UpperCamelCase  )


   @property
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      warnings.warn(
          "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."      ,    _UpperCamelCase      ,    )
      return self.image_processor_class






   @property
   def 						_lowerCamelCase       (			self  ):





      """simple docstring"""





      warnings.warn(
          "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."      ,    _UpperCamelCase      ,    )
      return self.image_processor






