




"""simple docstring"""







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


SCREAMING_SNAKE_CASE_	     =       logging.get_logger(__name__)

SCREAMING_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 				snake_case_			(					lowerCamelCase_			):







   """simple docstring"""




   def __init__(		self       ,					lowerCamelCase_=None       ,					lowerCamelCase_=None       ,					*lowerCamelCase_       ,					**lowerCamelCase_)				->   int:
       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__}'
           )
           UpperCamelCase              =	self.model.config
       else:
           UpperCamelCase              =	config

       UpperCamelCase              =	data_args
       UpperCamelCase              =	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:
           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   UpperCAmelCase__			(		self       ,					lowerCamelCase_)				->   Optional[int]:
       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=lowerCamelCase_       ,					optim=lowerCamelCase_       ,					**lowerCamelCase_       ,					)
           else:
               UpperCamelCase              =	optimizer_cls(lowerCamelCase_       ,					**lowerCamelCase_)

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







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_)				->   List[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=lowerCamelCase_)
       return scheduler







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







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_)				->   Any:
       UpperCamelCase              =	inputs.pop('''labels''')
       UpperCamelCase       ,	UpperCamelCase              =	self._compute_loss(lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_)
       return loss







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_       ,					lowerCamelCase_ = None       ,					)				->   Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
       UpperCamelCase              =	self._prepare_inputs(lowerCamelCase_)

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

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

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

       return (loss, logits, labels)







   def   UpperCAmelCase__			(		self       ,					lowerCamelCase_       ,					lowerCamelCase_)				->   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





import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor


class 							UpperCamelCase_       (					unittest.TestCase	):

    '''simple docstring'''






    def __init__(     self :      Tuple   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      int=13   ,					UpperCAmelCase__ :      Union[str, Any]=3   ,					UpperCAmelCase__ :      str=224   ,					UpperCAmelCase__ :      str=30   ,					UpperCAmelCase__ :      Tuple=400   ,					UpperCAmelCase__ :      int=True   ,					UpperCAmelCase__ :      Dict=None   ,					UpperCAmelCase__ :      List[Any]=True   ,					UpperCAmelCase__ :      Union[str, Any]=[0.5, 0.5, 0.5]   ,					UpperCAmelCase__ :      Tuple=[0.5, 0.5, 0.5]   ,					)		->str:





           '''simple docstring'''

           A__ 					=       size if size is not None else {'''height''': 18, '''width''': 18}
           A__ 					=       parent
           A__ 					=       batch_size
           A__ 					=       num_channels
           A__ 					=       image_size
           A__ 					=       min_resolution
           A__ 					=       max_resolution
           A__ 					=       do_resize
           A__ 					=       size
           A__ 					=       do_normalize
           A__ 					=       image_mean
           A__ 					=       image_std





    def 			SCREAMING_SNAKE_CASE				(     self :      Any)		->Optional[int]:





           '''simple docstring'''

           return {
               "image_mean": self.image_mean,
               "image_std": self.image_std,
               "do_normalize": self.do_normalize,
               "do_resize": self.do_resize,
               "size": self.size,
           }


@require_torch
@require_vision
class 							UpperCamelCase_       (					UpperCAmelCase__						,   unittest.TestCase	):

    '''simple docstring'''






    UpperCAmelCase__				    = ViTImageProcessor if is_vision_available() else None


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->str:





           '''simple docstring'''

           A__ 					=       EfficientFormerImageProcessorTester(self)


    @property
    def 			SCREAMING_SNAKE_CASE				(     self :      Dict)		->int:





           '''simple docstring'''

           return self.image_proc_tester.prepare_image_processor_dict()


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Dict:





           '''simple docstring'''

           A__ 					=       self.image_processing_class(**self.image_processor_dict)
           self.assertTrue(hasattr(UpperCAmelCase__   ,					'''image_mean'''))
           self.assertTrue(hasattr(UpperCAmelCase__   ,					'''image_std'''))
           self.assertTrue(hasattr(UpperCAmelCase__   ,					'''do_normalize'''))
           self.assertTrue(hasattr(UpperCAmelCase__   ,					'''do_resize'''))
           self.assertTrue(hasattr(UpperCAmelCase__   ,					'''size'''))


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Dict:





           '''simple docstring'''

           pass


    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       self.image_processing_class(**self.image_processor_dict)
           # create random PIL images
           A__ 					=       prepare_image_inputs(self.image_proc_tester   ,					equal_resolution=UpperCAmelCase__)
           for image in image_inputs:
                  self.assertIsInstance(UpperCAmelCase__   ,					Image.Image)

           # Test not batched input
           A__ 					=       image_processor(image_inputs[0]   ,					return_tensors='''pt''').pixel_values
           self.assertEqual(
               encoded_images.shape   ,					(
                   1,
                   self.image_proc_tester.num_channels,
                   self.image_proc_tester.size['''height'''],
                   self.image_proc_tester.size['''width'''],
               )   ,					)

           # Test batched
           A__ 					=       image_processor(UpperCAmelCase__   ,					return_tensors='''pt''').pixel_values
           self.assertEqual(
               encoded_images.shape   ,					(
                   self.image_proc_tester.batch_size,
                   self.image_proc_tester.num_channels,
                   self.image_proc_tester.size['''height'''],
                   self.image_proc_tester.size['''width'''],
               )   ,					)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->List[Any]:





           '''simple docstring'''

           A__ 					=       self.image_processing_class(**self.image_processor_dict)
           # create random numpy tensors
           A__ 					=       prepare_image_inputs(self.image_proc_tester   ,					equal_resolution=UpperCAmelCase__   ,					numpify=UpperCAmelCase__)
           for image in image_inputs:
                  self.assertIsInstance(UpperCAmelCase__   ,					np.ndarray)

           # Test not batched input
           A__ 					=       image_processor(image_inputs[0]   ,					return_tensors='''pt''').pixel_values
           self.assertEqual(
               encoded_images.shape   ,					(
                   1,
                   self.image_proc_tester.num_channels,
                   self.image_proc_tester.size['''height'''],
                   self.image_proc_tester.size['''width'''],
               )   ,					)

           # Test batched
           A__ 					=       image_processor(UpperCAmelCase__   ,					return_tensors='''pt''').pixel_values
           self.assertEqual(
               encoded_images.shape   ,					(
                   self.image_proc_tester.batch_size,
                   self.image_proc_tester.num_channels,
                   self.image_proc_tester.size['''height'''],
                   self.image_proc_tester.size['''width'''],
               )   ,					)





    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       self.image_processing_class(**self.image_processor_dict)
           # create random PyTorch tensors
           A__ 					=       prepare_image_inputs(self.image_proc_tester   ,					equal_resolution=UpperCAmelCase__   ,					torchify=UpperCAmelCase__)
           for image in image_inputs:
                  self.assertIsInstance(UpperCAmelCase__   ,					torch.Tensor)

           # Test not batched input
           A__ 					=       image_processor(image_inputs[0]   ,					return_tensors='''pt''').pixel_values
           self.assertEqual(
               encoded_images.shape   ,					(
                   1,
                   self.image_proc_tester.num_channels,
                   self.image_proc_tester.size['''height'''],
                   self.image_proc_tester.size['''width'''],
               )   ,					)

           # Test batched
           A__ 					=       image_processor(UpperCAmelCase__   ,					return_tensors='''pt''').pixel_values
           self.assertEqual(
               encoded_images.shape   ,					(
                   self.image_proc_tester.batch_size,
                   self.image_proc_tester.num_channels,
                   self.image_proc_tester.size['''height'''],
                   self.image_proc_tester.size['''width'''],
               )   ,					)


