





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     :     str							=	logging.get_logger(__name__)

UpperCAmelCase     :     str							=	{
    """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	(   UpperCamelCase__):







      def __init__(   self		,				lowerCAmelCase__=None		,				lowerCAmelCase__=None		,				*lowerCAmelCase__		,				**lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




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

        a__     : Union[str, Any]  =data_args
        a__     : str  =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:
          a__     : Tuple  =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__     : Any  =label_smoothed_nll_loss







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




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

        if self.lr_scheduler is None:
          a__     : Dict  =self._get_lr_scheduler(lowerCAmelCase__							)
        else:  # ignoring --lr_scheduler
          logger.warning("scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored."							)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Any:



        '''simple docstring'''




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







      def        _lowercase  (   self							)       ->						Optional[torch.utils.data.Sampler]:



        '''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		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[str]:



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







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : Any  =inputs.pop("labels"							)
        a__ ,							a__     : Any  =self._compute_loss(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)
        return loss







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__ = None		,				)       ->						Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:



        '''simple docstring'''




        a__     : List[str]  =self._prepare_inputs(lowerCAmelCase__							)

        a__     : Union[str, 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__     : Optional[int]  =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"]:
            a__     : Optional[int]  =self._pad_tensors_to_max_len(lowerCAmelCase__		,				gen_kwargs["max_length"]							)

        a__     : Tuple  =inputs.pop("labels"							)
        with torch.no_grad():
          # compute loss on predict data
          a__ ,							a__     : List[Any]  =self._compute_loss(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)

        a__     : List[str]  =loss.mean().detach()
        if self.args.prediction_loss_only:
          return (loss, None, None)

        a__     : Any  =generated_tokens if self.args.predict_with_generate else logits

        if labels.shape[-1] < gen_kwargs["max_length"]:
          a__     : List[str]  =self._pad_tensors_to_max_len(lowerCAmelCase__		,				gen_kwargs["max_length"]							)

        return (loss, logits, labels)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : 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__     : List[Any]  =pad_token_id * torch.ones(
            (tensor.shape[0], max_length)		,				dtype=tensor.dtype		,				device=tensor.device							)
        a__     : List[str]  =tensor
        return padded_tensor













import warnings

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase     :     Dict							=	logging.get_logger(__name__)

UpperCAmelCase     :     Dict							=	{
    """RUCAIBox/mvp""": """https://huggingface.co/RUCAIBox/mvp/resolve/main/config.json""",
}







class     __lowerCAmelCase	(   UpperCamelCase__):
      _lowercase    : List[Any]            =      """mvp"""
      _lowercase    : Tuple            =      ["""past_key_values"""]
      _lowercase    : Dict            =      {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}







      def __init__(   self		,				lowerCAmelCase__=5_0_2_6_7		,				lowerCAmelCase__=1_0_2_4		,				lowerCAmelCase__=1_2		,				lowerCAmelCase__=4_0_9_6		,				lowerCAmelCase__=1_6		,				lowerCAmelCase__=1_2		,				lowerCAmelCase__=4_0_9_6		,				lowerCAmelCase__=1_6		,				lowerCAmelCase__=0.0		,				lowerCAmelCase__=0.0		,				lowerCAmelCase__="gelu"		,				lowerCAmelCase__=1_0_2_4		,				lowerCAmelCase__=0.1		,				lowerCAmelCase__=0.0		,				lowerCAmelCase__=0.0		,				lowerCAmelCase__=0.02		,				lowerCAmelCase__=0.0		,				lowerCAmelCase__=False		,				lowerCAmelCase__=True		,				lowerCAmelCase__=1		,				lowerCAmelCase__=0		,				lowerCAmelCase__=2		,				lowerCAmelCase__=True		,				lowerCAmelCase__=2		,				lowerCAmelCase__=2		,				lowerCAmelCase__=False		,				lowerCAmelCase__=1_0_0		,				lowerCAmelCase__=8_0_0		,				**lowerCAmelCase__		,				)       ->						str:



        '''simple docstring'''




        a__     : Dict  =vocab_size
        a__     : List[str]  =max_position_embeddings
        a__     : List[str]  =d_model
        a__     : Optional[Any]  =encoder_ffn_dim
        a__     : Dict  =encoder_layers
        a__     : List[str]  =encoder_attention_heads
        a__     : List[str]  =decoder_ffn_dim
        a__     : Tuple  =decoder_layers
        a__     : Tuple  =decoder_attention_heads
        a__     : Any  =dropout
        a__     : str  =attention_dropout
        a__     : str  =activation_dropout
        a__     : Optional[int]  =activation_function
        a__     : Union[str, Any]  =init_std
        a__     : Dict  =encoder_layerdrop
        a__     : List[str]  =decoder_layerdrop
        a__     : Union[str, Any]  =classifier_dropout
        a__     : Union[str, Any]  =use_cache
        a__     : Optional[Any]  =encoder_layers
        a__     : List[str]  =scale_embedding  # scale factor will be sqrt(d_model) if True
        a__     : int  =use_prompt
        a__     : Union[str, Any]  =prompt_length
        a__     : Dict  =prompt_mid_dim

        super().__init__(
            pad_token_id=lowerCAmelCase__		,				bos_token_id=lowerCAmelCase__		,				eos_token_id=lowerCAmelCase__		,				is_encoder_decoder=lowerCAmelCase__		,				decoder_start_token_id=lowerCAmelCase__		,				forced_eos_token_id=lowerCAmelCase__		,				**lowerCAmelCase__		,				)

        if self.forced_bos_token_id is None and kwargs.get("force_bos_token_to_be_generated"		,				lowerCAmelCase__							):
          a__     : Tuple  =self.bos_token_id
          warnings.warn(
              F'''Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. '''
              "The config can simply be saved and uploaded again to be fixed."							)






