





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       :   Union[str, Any]				       =							logging.get_logger(__name__)

__UpperCAmelCase       :   Dict				       =							{
    "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       (							__lowerCamelCase			):





       '''simple docstring'''







       def __init__(  self :    str							,				A :    int=None							,				A :    List[str]=None							,				*A :    str							,				**A :    List[str]		):
         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__}'''
           )
           __snake_case:						Any   =						self.model.config
         else:
           __snake_case:						Union[str, Any]   =						config

         __snake_case:						Any   =						data_args
         __snake_case:						List[Any]   =						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:
           __snake_case:						int   =						torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id		)
         else:
           # dynamically import label_smoothed_nll_loss
           from utils import label_smoothed_nll_loss

           __snake_case:						List[Any]   =						label_smoothed_nll_loss




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

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




       def 							UpperCAmelCase__ (  self :    Union[str, Any]							,				A :    str		):
         __snake_case:						Dict   =						arg_to_scheduler[self.args.lr_scheduler]
         if self.args.lr_scheduler == "constant":
           __snake_case:						Optional[int]   =						schedule_func(self.optimizer		)
         elif self.args.lr_scheduler == "constant_w_warmup":
           __snake_case:						List[Any]   =						schedule_func(self.optimizer							,				num_warmup_steps=self.args.warmup_steps		)
         else:
           __snake_case:						Union[str, Any]   =						schedule_func(
               self.optimizer							,				num_warmup_steps=self.args.warmup_steps							,				num_training_steps=A		)
         return scheduler




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




       def 							UpperCAmelCase__ (  self :    Tuple							,				A :    List[str]							,				A :    Union[str, Any]		):
         __snake_case:						Any   =						inputs.pop("""labels"""		)
         __snake_case  ,   __snake_case:						Tuple   =						self._compute_loss(A							,				A							,				A		)
         return loss




       def 							UpperCAmelCase__ (  self :    int							,				A :    Union[str, Any]							,				A :    Union[str, Any]							,				A :    List[Any]							,				A :    List[Any] = None							,				):
         __snake_case:						Tuple   =						self._prepare_inputs(A		)

         __snake_case:						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:
           __snake_case:						Tuple   =						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"]:
             __snake_case:						int   =						self._pad_tensors_to_max_len(A							,				gen_kwargs["""max_length"""]		)

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

         __snake_case:						Optional[Any]   =						loss.mean().detach()
         if self.args.prediction_loss_only:
           return (loss, None, None)

         __snake_case:						Optional[int]   =						generated_tokens if self.args.predict_with_generate else logits

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

         return (loss, logits, labels)






       def 							UpperCAmelCase__ (  self :    Union[str, Any]							,				A :    Optional[Any]							,				A :    int		):
         # If PAD token is not defined at least EOS token has to be defined
         __snake_case:						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}'''		)

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










import comet  # From: unbabel-comet
import torch

import datasets


UpperCAmelCase__   	=							datasets.logging.get_logger(__name__)

UpperCAmelCase__   	=							'''\
@inproceedings{rei-EtAl:2020:WMT,
   author    = {Rei, Ricardo  and  Stewart, Craig  and  Farinha, Ana C  and  Lavie, Alon},
   title     = {Unbabel\'s Participation in the WMT20 Metrics Shared Task},
   booktitle      = {Proceedings of the Fifth Conference on Machine Translation},
   month          = {November},
   year           = {2020},
   address        = {Online},
   publisher      = {Association for Computational Linguistics},
   pages     = {909--918},
}
@inproceedings{rei-etal-2020-comet,
   title = "{COMET}: A Neural Framework for {MT} Evaluation",
   author = "Rei, Ricardo  and
      Stewart, Craig  and
      Farinha, Ana C  and
      Lavie, Alon",
   booktitle = "Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)",
   month = nov,
   year = "2020",
   address = "Online",
   publisher = "Association for Computational Linguistics",
   url = "https://www.aclweb.org/anthology/2020.emnlp-main.213",
   pages = "2685--2702",
}
'''

UpperCAmelCase__   	=							'''\
Crosslingual Optimized Metric for Evaluation of Translation (COMET) is an open-source framework used to train Machine Translation metrics that achieve high levels of correlation with different types of human judgments (HTER, DA\'s or MQM).
With the release of the framework the authors also released fully trained models that were used to compete in the WMT20 Metrics Shared Task achieving SOTA in that years competition.

See the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.
'''

UpperCAmelCase__   	=							'''
COMET score.

Args:

`sources` (list of str): Source sentences
`predictions` (list of str): candidate translations
`references` (list of str): reference translations
`cuda` (bool): If set to True, runs COMET using GPU
`show_progress` (bool): Shows progress
`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.

Returns:
    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.
    `scores`: List of scores.

Examples:

    >>> comet_metric = datasets.load_metric(\'comet\')
    >>> # comet_metric = load_metric(\'comet\', \'wmt20-comet-da\')  # you can also choose which model to use
    >>> source = ["Dem Feuer konnte Einhalt geboten werden", "Schulen und Kindergärten wurden eröffnet."]
    >>> hypothesis = ["The fire could be stopped", "Schools and kindergartens were open"]
    >>> reference = ["They were able to control the fire.", "Schools and kindergartens opened"]
    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)
    >>> print([round(v, 2) for v in results["scores"]])
    [0.19, 0.92]
'''







@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION   ,							_KWARGS_DESCRIPTION)
class       lowerCamelCase__   (       datasets.Metric):





     def      __A					(self       )     ->   Optional[int]:
      return datasets.MetricInfo(
          description=_DESCRIPTION		,     citation=_CITATION		,     homepage='''https://unbabel.github.io/COMET/html/index.html'''		,     inputs_description=_KWARGS_DESCRIPTION		,     features=datasets.Features(
              {
                  '''sources''': datasets.Value('''string'''		,     id='''sequence'''       ),
                  '''predictions''': datasets.Value('''string'''		,     id='''sequence'''       ),
                  '''references''': datasets.Value('''string'''		,     id='''sequence'''       ),
              }       )		,     codebase_urls=['''https://github.com/Unbabel/COMET''']		,     reference_urls=[
              '''https://github.com/Unbabel/COMET''',
              '''https://www.aclweb.org/anthology/2020.emnlp-main.213/''',
              '''http://www.statmt.org/wmt20/pdf/2020.wmt-1.101.pdf6''',
          ]		,     )





     def      __A					(self		,     UpperCAmelCase       )     ->   Dict:
      if self.config_name == "default":
       _lowercase							      =comet.load_from_checkpoint(comet.download_model('''wmt20-comet-da'''       )       )
      else:
       _lowercase							      =comet.load_from_checkpoint(comet.download_model(self.config_name       )       )





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase=None		,     UpperCAmelCase=False       )     ->   int:
      if gpus is None:
       _lowercase							      =1 if torch.cuda.is_available() else 0
      _lowercase							      ={'''src''': sources, '''mt''': predictions, '''ref''': references}
      _lowercase							      =[dict(zip(UpperCAmelCase		,     UpperCAmelCase       )       ) for t in zip(*data.values()       )]
      _lowercase      ,    _lowercase							      =self.scorer.predict(UpperCAmelCase		,     gpus=UpperCAmelCase		,     progress_bar=UpperCAmelCase       )
      return {"mean_score": mean_score, "scores": scores}

