




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_ : List[Any]    		=       logging.get_logger(__name__)

snake_case_ : 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 __snake_case			(			a		):







   def __init__(      self		:						Optional[Any]					,						_snake_case		:						int=None					,						_snake_case		:						Optional[Any]=None					,						*_snake_case		:						Any					,						**_snake_case		:						Optional[Any]):

        """simple docstring"""


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

        UpperCAmelCase_          =							data_args
        UpperCAmelCase_          =							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:
             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 lowerCamelCase						(      self		:						Any					,						_snake_case		:						int):

        """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=_snake_case					,						optim=_snake_case					,						**_snake_case					,						)
             else:
                  UpperCAmelCase_          =							optimizer_cls(_snake_case					,						**_snake_case)

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







   def lowerCamelCase						(      self		:						str					,						_snake_case		:						Union[str, Any]):

        """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=_snake_case)
        return scheduler







   def lowerCamelCase						(      self		:						Dict):

        """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		:						str					,						_snake_case		:						Optional[int]					,						_snake_case		:						int					,						_snake_case		:						Optional[Any]):

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







   def lowerCamelCase						(      self		:						List[str]					,						_snake_case		:						Union[str, Any]					,						_snake_case		:						Dict):

        """simple docstring"""


        UpperCAmelCase_          =							inputs.pop('''labels''')
        UpperCAmelCase_     ,       UpperCAmelCase_          =							self._compute_loss(_snake_case					,						_snake_case					,						_snake_case)
        return loss







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

        """simple docstring"""


        UpperCAmelCase_          =							self._prepare_inputs(_snake_case)

        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''']					,						**_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"]:
                  UpperCAmelCase_          =							self._pad_tensors_to_max_len(_snake_case					,						gen_kwargs['''max_length'''])

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

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

        return (loss, logits, labels)







   def lowerCamelCase						(      self		:						str					,						_snake_case		:						Optional[Any]					,						_snake_case		:						Optional[Any]):

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














"""simple docstring"""






from typing import TYPE_CHECKING

from ...file_utils import _LazyModule, is_tokenizers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


UpperCAmelCase_   : str   			= {"""configuration_gpt_neox""": ["""GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP""", """GPTNeoXConfig"""]}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_   : Dict   			= ["""GPTNeoXTokenizerFast"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_   : List[str]   			= [
      """GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """GPTNeoXForCausalLM""",
      """GPTNeoXForQuestionAnswering""",
      """GPTNeoXForSequenceClassification""",
      """GPTNeoXForTokenClassification""",
      """GPTNeoXLayer""",
      """GPTNeoXModel""",
      """GPTNeoXPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_gpt_neox import GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXConfig

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_gpt_neox_fast import GPTNeoXTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_gpt_neox import (
        GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST,
        GPTNeoXForCausalLM,
        GPTNeoXForQuestionAnswering,
        GPTNeoXForSequenceClassification,
        GPTNeoXForTokenClassification,
        GPTNeoXLayer,
        GPTNeoXModel,
        GPTNeoXPreTrainedModel,
    )


else:
  import sys

  UpperCAmelCase_   : Tuple   			= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


