




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


UpperCamelCase__ :   Optional[int]						=							logging.get_logger(__name__)

UpperCamelCase__ :   Tuple						=							{
    '''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_       (					lowerCamelCase_							):






				def __init__(			self      ,snake_case__=None      ,snake_case__=None      ,*snake_case__      ,**snake_case__ ):
							super().__init__(*snake_case__      ,**snake_case__ )

							if config is None:
										assert isinstance(self.model      ,snake_case__ ), (
										    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
										    F' {self.model.__class__}'
										)
										SCREAMING_SNAKE_CASE_   : Any					  =      self.model.config
							else:
										SCREAMING_SNAKE_CASE_   : Dict					  =      config

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

										SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      label_smoothed_nll_loss






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

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






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






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






				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      inputs.pop('labels' )
							SCREAMING_SNAKE_CASE_  ,    SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      self._compute_loss(snake_case__      ,snake_case__      ,snake_case__ )
							return loss






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__ = None      ,):
							SCREAMING_SNAKE_CASE_   : Tuple					  =      self._prepare_inputs(snake_case__ )

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

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

							SCREAMING_SNAKE_CASE_   : str					  =      loss.mean().detach()
							if self.args.prediction_loss_only:
										return (loss, None, None)

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

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

							return (loss, logits, labels)






				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							# If PAD token is not defined at least EOS token has to be defined
							SCREAMING_SNAKE_CASE_   : Tuple					  =      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}' )

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









import os
import sys
import unittest


UpperCamelCase__ :   Optional[Any]						=							os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, '''utils'''))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
UpperCamelCase__ :   Tuple						=							os.path.join(git_repo_path, '''src''', '''transformers''')

UpperCamelCase__ :   List[Any]						=							'''
{0} = None
'''

UpperCamelCase__ :   str						=							'''
class {0}(metaclass=DummyObject):
    _backends = {1}

    def __init__(self, *args, **kwargs):
        requires_backends(self, {1})
'''


UpperCamelCase__ :   List[Any]						=							'''
def {0}(*args, **kwargs):
    requires_backends({0}, {1})
'''







class 				lowerCAmelCase_       (					unittest.TestCase							):






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : int					  =      find_backend('    _import_structure["models.albert"].append("AlbertTokenizerFast")' )
							self.assertIsNone(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      find_backend('    if not is_tokenizers_available():' )
							self.assertEqual(snake_case__      ,'tokenizers' )

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      find_backend('    if not is_tensorflow_text_available():' )
							self.assertEqual(snake_case__      ,'tensorflow_text' )

							SCREAMING_SNAKE_CASE_   : int					  =      find_backend('    if not (is_sentencepiece_available() and is_tokenizers_available()):' )
							self.assertEqual(snake_case__      ,'sentencepiece_and_tokenizers' )

							SCREAMING_SNAKE_CASE_   : str					  =      find_backend(
							    '    if not (is_sentencepiece_available() and is_tensorflow_text_available()):' )
							self.assertEqual(snake_case__      ,'sentencepiece_and_tensorflow_text' )

							SCREAMING_SNAKE_CASE_   : Tuple					  =      find_backend(
							    '    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):' )
							self.assertEqual(snake_case__      ,'sentencepiece_and_tokenizers_and_vision' )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      read_init()
							# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
							self.assertIn('torch'      ,snake_case__ )
							self.assertIn('tensorflow_text'      ,snake_case__ )
							self.assertIn('sentencepiece_and_tokenizers'      ,snake_case__ )

							# Likewise, we can't assert on the exact content of a key
							self.assertIn('BertModel'      ,objects['torch'] )
							self.assertIn('TFBertModel'      ,objects['tf'] )
							self.assertIn('FlaxBertModel'      ,objects['flax'] )
							self.assertIn('BertModel'      ,objects['torch'] )
							self.assertIn('TFBertTokenizer'      ,objects['tensorflow_text'] )
							self.assertIn('convert_slow_tokenizer'      ,objects['sentencepiece_and_tokenizers'] )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      create_dummy_object('CONSTANT'      ,'\'torch\'' )
							self.assertEqual(snake_case__      ,'\nCONSTANT = None\n' )

							SCREAMING_SNAKE_CASE_   : Dict					  =      create_dummy_object('function'      ,'\'torch\'' )
							self.assertEqual(
							    snake_case__      ,'\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n' )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      '\nclass FakeClass(metaclass=DummyObject):\n    _backends = \'torch\'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, \'torch\')\n'
							SCREAMING_SNAKE_CASE_   : Dict					  =      create_dummy_object('FakeClass'      ,'\'torch\'' )
							self.assertEqual(snake_case__      ,snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      '# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, ["torch"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = ["torch"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, ["torch"])\n'
							SCREAMING_SNAKE_CASE_   : Dict					  =      create_dummy_files({'torch': ['CONSTANT', 'function', 'FakeClass']} )
							self.assertEqual(dummy_files['torch']      ,snake_case__ )



