






from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
							import PIL


A_       :Any       =						logging.get_logger(__name__)




class 				__A     (			a ):





				"""simple docstring"""







				UpperCamelCase__ :      int			 =["""pixel_values"""]


				def __init__( self							,  lowerCamelCase__ = True							,  lowerCamelCase__ = None							,  lowerCamelCase__ = PIL.Image.BICUBIC							,  lowerCamelCase__ = True							,  lowerCamelCase__ = None							,  lowerCamelCase__ = 1 / 255							,  lowerCamelCase__ = True							,  lowerCamelCase__ = True							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  **lowerCamelCase__							,  ):




									"""simple docstring"""

									super().__init__(**lowerCamelCase__    )
									__UpperCamelCase      :					Optional[Any]       		=size if size is not None else {'height': 256, 'width': 256}
									__UpperCamelCase      :					List[str]       		=get_size_dict(lowerCamelCase__    )
									__UpperCamelCase      :					Tuple       		=crop_size if crop_size is not None else {'height': 224, 'width': 224}
									__UpperCamelCase      :					Optional[int]       		=get_size_dict(lowerCamelCase__							,  param_name='crop_size'    )

									__UpperCamelCase      :					Tuple       		=do_resize
									__UpperCamelCase      :					Any       		=size
									__UpperCamelCase      :					Dict       		=resample
									__UpperCamelCase      :					Tuple       		=do_center_crop
									__UpperCamelCase      :					int       		=crop_size
									__UpperCamelCase      :					Optional[int]       		=do_rescale
									__UpperCamelCase      :					List[Any]       		=rescale_factor
									__UpperCamelCase      :					List[Any]       		=do_normalize
									__UpperCamelCase      :					Union[str, Any]       		=image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
									__UpperCamelCase      :					Optional[Any]       		=image_std if image_std is not None else IMAGENET_STANDARD_STD


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__ = PIL.Image.BICUBIC							,  lowerCamelCase__ = None							,  **lowerCamelCase__							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		=get_size_dict(lowerCamelCase__    )
									if "height" not in size or "width" not in size:
														raise ValueError(f'The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'    )
									return resize(
									    lowerCamelCase__							,  size=(size['height'], size['width'])							,  resample=lowerCamelCase__							,  data_format=lowerCamelCase__							,  **lowerCamelCase__    )


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__ = None							,  **lowerCamelCase__							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[int]       		=get_size_dict(lowerCamelCase__    )
									if "height" not in size or "width" not in size:
														raise ValueError(f'The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'    )
									return center_crop(lowerCamelCase__							,  size=(size['height'], size['width'])							,  data_format=lowerCamelCase__							,  **lowerCamelCase__    )


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__ = None							,  **lowerCamelCase__							,  ):




									"""simple docstring"""

									return rescale(lowerCamelCase__							,  scale=lowerCamelCase__							,  data_format=lowerCamelCase__							,  **lowerCamelCase__    )


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__ = None							,  **lowerCamelCase__							,  ):




									"""simple docstring"""

									return normalize(lowerCamelCase__							,  mean=lowerCamelCase__							,  std=lowerCamelCase__							,  data_format=lowerCamelCase__							,  **lowerCamelCase__    )







				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__=None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = None							,  lowerCamelCase__ = ChannelDimension.FIRST							,  **lowerCamelCase__							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[Any]       		=do_resize if do_resize is not None else self.do_resize
									__UpperCamelCase      :					Union[str, Any]       		=resample if resample is not None else self.resample
									__UpperCamelCase      :					Optional[Any]       		=do_center_crop if do_center_crop is not None else self.do_center_crop
									__UpperCamelCase      :					Any       		=do_rescale if do_rescale is not None else self.do_rescale
									__UpperCamelCase      :					Dict       		=rescale_factor if rescale_factor is not None else self.rescale_factor
									__UpperCamelCase      :					Optional[Any]       		=do_normalize if do_normalize is not None else self.do_normalize
									__UpperCamelCase      :					Optional[Any]       		=image_mean if image_mean is not None else self.image_mean
									__UpperCamelCase      :					Any       		=image_std if image_std is not None else self.image_std

									__UpperCamelCase      :					Union[str, Any]       		=size if size is not None else self.size
									__UpperCamelCase      :					List[str]       		=get_size_dict(lowerCamelCase__    )
									__UpperCamelCase      :					Tuple       		=crop_size if crop_size is not None else self.crop_size
									__UpperCamelCase      :					Tuple       		=get_size_dict(lowerCamelCase__							,  param_name='crop_size'    )

									__UpperCamelCase      :					List[Any]       		=make_list_of_images(lowerCamelCase__    )

									if not valid_images(lowerCamelCase__    ):
														raise ValueError(
														    'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
														    'torch.Tensor, tf.Tensor or jax.ndarray.'    )

									if do_resize and size is None or resample is None:
														raise ValueError('Size and resample must be specified if do_resize is True.'    )

									if do_center_crop and crop_size is None:
														raise ValueError('Crop size must be specified if do_center_crop is True.'    )

									if do_rescale and rescale_factor is None:
														raise ValueError('Rescale factor must be specified if do_rescale is True.'    )

									if do_normalize and (image_mean is None or image_std is None):
														raise ValueError('Image mean and std must be specified if do_normalize is True.'    )

									# All transformations expect numpy arrays.
									__UpperCamelCase      :					Optional[int]       		=[to_numpy_array(lowerCamelCase__    ) for image in images]

									if do_resize:
														__UpperCamelCase      :					Optional[Any]       		=[self.resize(image=lowerCamelCase__							,  size=lowerCamelCase__							,  resample=lowerCamelCase__    ) for image in images]

									if do_center_crop:
														__UpperCamelCase      :					int       		=[self.center_crop(image=lowerCamelCase__							,  size=lowerCamelCase__    ) for image in images]

									if do_rescale:
														__UpperCamelCase      :					int       		=[self.rescale(image=lowerCamelCase__							,  scale=lowerCamelCase__    ) for image in images]

									if do_normalize:
														__UpperCamelCase      :					Optional[int]       		=[self.normalize(image=lowerCamelCase__							,  mean=lowerCamelCase__							,  std=lowerCamelCase__    ) for image in images]

									__UpperCamelCase      :					Optional[Any]       		=[to_channel_dimension_format(lowerCamelCase__							,  lowerCamelCase__    ) for image in images]

									__UpperCamelCase      :					List[str]       		={'pixel_values': images}
									return BatchFeature(data=lowerCamelCase__							,  tensor_type=lowerCamelCase__    )



'''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


__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 UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  def __init__(		self      :     Optional[int]							,_a      :     Optional[Any]=None							,_a      :     Dict=None							,*_a      :     int							,**_a      :     str       ):


     '''simple docstring'''






     super().__init__(*_a							,**_a       )

     if config is None:
        assert isinstance(self.model							,_a       ), (
            "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
            F""" {self.model.__class__}"""
        )
        _a				:	List[Any]			=   self.model.config
     else:
        _a				:	Optional[int]			=   config

     _a				:	List[str]			=   data_args
     _a				:	List[Any]			=   self.config.tgt_vocab_size if isinstance(self.config							,_a       ) 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:
        _a				:	List[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

        _a				:	Tuple			=   label_smoothed_nll_loss


  def   __lowercase		(		self      :     List[str]							,_a      :     int       ):


     '''simple docstring'''






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

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


  def   __lowercase		(		self      :     List[Any]							,_a      :     List[Any]       ):


     '''simple docstring'''






     _a				:	str			=   arg_to_scheduler[self.args.lr_scheduler]
     if self.args.lr_scheduler == "constant":
        _a				:	int			=   schedule_func(self.optimizer       )
     elif self.args.lr_scheduler == "constant_w_warmup":
        _a				:	List[str]			=   schedule_func(self.optimizer							,num_warmup_steps=self.args.warmup_steps       )
     else:
        _a				:	Optional[int]			=   schedule_func(
            self.optimizer							,num_warmup_steps=self.args.warmup_steps							,num_training_steps=_a       )
     return scheduler


  def   __lowercase		(		self      :     Tuple       ):


     '''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   __lowercase		(		self      :     Dict							,_a      :     Dict							,_a      :     Any							,_a      :     Dict       ):


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


  def   __lowercase		(		self      :     Optional[int]							,_a      :     Union[str, Any]							,_a      :     List[Any]       ):


     '''simple docstring'''






     _a				:	Optional[int]			=   inputs.pop('labels'       )
     _a,  _a				:	int			=   self._compute_loss(_a							,_a							,_a       )
     return loss


  def   __lowercase		(		self      :     Optional[Any]							,_a      :     nn.Module							,_a      :     Dict[str, Union[torch.Tensor, Any]]							,_a      :     bool							,_a      :     Optional[List[str]] = None							,):


     '''simple docstring'''






     _a				:	int			=   self._prepare_inputs(_a       )

     _a				:	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:
        _a				:	int			=   self.model.generate(
            inputs['input_ids']							,attention_mask=inputs['attention_mask']							,**_a							,)
        # in case the batch is shorter than max length, the output should be padded
        if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
           _a				:	int			=   self._pad_tensors_to_max_len(_a							,gen_kwargs['max_length']       )

     _a				:	Union[str, Any]			=   inputs.pop('labels'       )
     with torch.no_grad():
        # compute loss on predict data
        _a,  _a				:	Optional[int]			=   self._compute_loss(_a							,_a							,_a       )

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

     _a				:	Optional[Any]			=   generated_tokens if self.args.predict_with_generate else logits

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

     return (loss, logits, labels)


  def   __lowercase		(		self      :     str							,_a      :     Tuple							,_a      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[Any]			=   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}"""       )

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



