






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

snake_case    				=		{
    """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  SCREAMING_SNAKE_CASE			(			lowerCAmelCase					):
     '''simple docstring'''


     def __init__(				self       :	Dict					, UpperCAmelCase_       :	Optional[Any]=None					, UpperCAmelCase_       :	Optional[int]=None					, *UpperCAmelCase_       :	Dict					, **UpperCAmelCase_       :	Tuple      ):
       super().__init__(*UpperCAmelCase_					, **UpperCAmelCase_      )

       if config is None:
         assert isinstance(self.model					, UpperCAmelCase_      ), (
             "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       :      List[str]							=    self.model.config
       else:
         SCREAMING_SNAKE_CASE       :      Tuple							=    config

       SCREAMING_SNAKE_CASE       :      Any							=    data_args
       SCREAMING_SNAKE_CASE       :      Tuple							=    self.config.tgt_vocab_size if isinstance(self.config					, UpperCAmelCase_      ) 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       :      Union[str, 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       :      Dict							=    label_smoothed_nll_loss


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

       if self.lr_scheduler is None:
         SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    self._get_lr_scheduler(UpperCAmelCase_      )
       else:  # ignoring --lr_scheduler
         logger.warning("scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored."      )


     def 				_A       (				self       :	Optional[Any]					, UpperCAmelCase_       :	Any      ):
       SCREAMING_SNAKE_CASE       :      Dict							=    arg_to_scheduler[self.args.lr_scheduler]
       if self.args.lr_scheduler == "constant":
         SCREAMING_SNAKE_CASE       :      List[str]							=    schedule_func(self.optimizer      )
       elif self.args.lr_scheduler == "constant_w_warmup":
         SCREAMING_SNAKE_CASE       :      Dict							=    schedule_func(self.optimizer					, num_warmup_steps=self.args.warmup_steps      )
       else:
         SCREAMING_SNAKE_CASE       :      Tuple							=    schedule_func(
             self.optimizer					, num_warmup_steps=self.args.warmup_steps					, num_training_steps=UpperCAmelCase_      )
       return scheduler


     def 				_A       (				self       :	Union[str, Any]      ):
       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 				_A       (				self       :	List[Any]					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	List[Any]					, UpperCAmelCase_       :	Tuple      ):
       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       :      Any							=    model(**UpperCAmelCase_					, use_cache=UpperCAmelCase_      )[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       :      str							=    model(**UpperCAmelCase_					, labels=UpperCAmelCase_					, use_cache=UpperCAmelCase_      )[:2]
       else:
         # compute label smoothed loss
         SCREAMING_SNAKE_CASE       :      Dict							=    model(**UpperCAmelCase_					, use_cache=UpperCAmelCase_      )[0]
         SCREAMING_SNAKE_CASE       :      Dict							=    torch.nn.functional.log_softmax(UpperCAmelCase_					, dim=-1      )
         SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE       :      Any							=    self.loss_fn(UpperCAmelCase_					, UpperCAmelCase_					, self.args.label_smoothing					, ignore_index=self.config.pad_token_id      )
       return loss, logits


     def 				_A       (				self       :	Optional[int]					, UpperCAmelCase_       :	str					, UpperCAmelCase_       :	List[str]      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    inputs.pop("labels"      )
       SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE       :      List[Any]							=    self._compute_loss(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_      )
       return loss


     def 				_A       (				self       :	Tuple					, UpperCAmelCase_       :	nn.Module					, UpperCAmelCase_       :	Dict[str, Union[torch.Tensor, Any]]					, UpperCAmelCase_       :	bool					, UpperCAmelCase_       :	Optional[List[str]] = None					, ):
       SCREAMING_SNAKE_CASE       :      List[Any]							=    self._prepare_inputs(UpperCAmelCase_      )

       SCREAMING_SNAKE_CASE       :      Optional[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:
         SCREAMING_SNAKE_CASE       :      Optional[int]							=    self.model.generate(
             inputs["input_ids"]					, attention_mask=inputs["attention_mask"]					, **UpperCAmelCase_					, )
         # 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       :      List[str]							=    self._pad_tensors_to_max_len(UpperCAmelCase_					, gen_kwargs["max_length"]      )

       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    inputs.pop("labels"      )
       with torch.no_grad():
         # compute loss on predict data
         SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE       :      Dict							=    self._compute_loss(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_      )

       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    loss.mean().detach()
       if self.args.prediction_loss_only:
         return (loss, None, None)

       SCREAMING_SNAKE_CASE       :      List[str]							=    generated_tokens if self.args.predict_with_generate else logits

       if labels.shape[-1] < gen_kwargs["max_length"]:
         SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    self._pad_tensors_to_max_len(UpperCAmelCase_					, gen_kwargs["max_length"]      )

       return (loss, logits, labels)






     def 				_A       (				self       :	Tuple					, UpperCAmelCase_       :	Dict					, UpperCAmelCase_       :	Optional[int]      ):
       # If PAD token is not defined at least EOS token has to be defined
       SCREAMING_SNAKE_CASE       :      List[str]							=    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       :      Optional[int]							=    pad_token_id * torch.ones(
           (tensor.shape[0], max_length)					, dtype=tensor.dtype					, device=tensor.device      )
       SCREAMING_SNAKE_CASE       :      int							=    tensor
       return padded_tensor














# Copyright 2022 The HuggingFace Team and The OpenBMB Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


snake_case    				=		{
    """configuration_cpmant""": ["""CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """CpmAntConfig"""],
    """tokenization_cpmant""": ["""CpmAntTokenizer"""],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   snake_case    				=		[
       """CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """CpmAntForCausalLM""",
       """CpmAntModel""",
       """CpmAntPreTrainedModel""",
   ]


if TYPE_CHECKING:
   from .configuration_cpmant import CPMANT_PRETRAINED_CONFIG_ARCHIVE_MAP, CpmAntConfig
   from .tokenization_cpmant import CpmAntTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_cpmant import (
          CPMANT_PRETRAINED_MODEL_ARCHIVE_LIST,
          CpmAntForCausalLM,
          CpmAntModel,
          CpmAntPreTrainedModel,
      )


else:
   import sys

   snake_case    				=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






