





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


lowercase    : str    	=  logging.get_logger(__name__)

lowercase    : Optional[Any]    	=  {
    """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,       A=None,       A=None,       *A,       **A      ):





												"""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__}'''
																			)
																			_UpperCamelCase              =						self.model.config
												else:
																			_UpperCamelCase              =						config

												_UpperCamelCase              =						data_args
												_UpperCamelCase              =						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:
																			_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,       A      ):





												"""simple docstring"""

												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=A,       optim=A,       **A,       )
																			else:
																										_UpperCamelCase              =						optimizer_cls(A,       **A      )

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

					def 	_UpperCamelCase     (      self,       A      ):





												"""simple docstring"""

												_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=A      )
												return scheduler

					def 	_UpperCamelCase     (      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 	_UpperCamelCase     (      self,       A,       A,       A      ):





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

					def 	_UpperCamelCase     (      self,       A,       A      ):





												"""simple docstring"""

												_UpperCamelCase              =						inputs.pop('''labels'''      )
												_UpperCamelCase       ,		_UpperCamelCase              =						self._compute_loss(A,       A,       A      )
												return loss

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





												"""simple docstring"""

												_UpperCamelCase              =						self._prepare_inputs(A      )

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

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

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

												return (loss, logits, labels)



					def 	_UpperCamelCase     (      self,       A,       A      ):





												"""simple docstring"""

												_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 warnings

from ...utils import logging
from .image_processing_clip import CLIPImageProcessor


lowercase    : str    	=  logging.get_logger(__name__)

class  __A(						__UpperCAmelCase ):

     def __init__(      self,       *A,       **A      ):





            """simple docstring"""

            warnings.warn(
                '''The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please'''
                ''' use CLIPImageProcessor instead.''',       A,       )
            super().__init__(*A,       **A      )






