"""simple docstring"""


import functools
from contextlib import contextmanager

import bert_score
from packaging import version

import datasets







@contextmanager
def       UpperCAmelCase					(	)       ->   str:
     def filter_log(UpperCAmelCase      ):
          return False if "This IS expected if you are initializing" in record.msg else True

     snake_case_       =      datasets.utils.logging.get_logger('transformers.modeling_utils'      )
     logger.addFilter(UpperCAmelCase      )
     try:
          yield
     finally:
          logger.removeFilter(UpperCAmelCase      )


__UpperCamelCase								=    '''\
@inproceedings{bert-score,
  title={BERTScore: Evaluating Text Generation with BERT},
  author={Tianyi Zhang* and Varsha Kishore* and Felix Wu* and Kilian Q. Weinberger and Yoav Artzi},
  booktitle={International Conference on Learning Representations},
  year={2020},
  url={https://openreview.net/forum?id=SkeHuCVFDr}
}
'''

__UpperCamelCase								=    '''\
BERTScore leverages the pre-trained contextual embeddings from BERT and matches words in candidate and reference
sentences by cosine similarity.
It has been shown to correlate with human judgment on sentence-level and system-level evaluation.
Moreover, BERTScore computes precision, recall, and F1 measure, which can be useful for evaluating different language
generation tasks.

See the project\'s README at https://github.com/Tiiiger/bert_score#readme for more information.
'''

__UpperCamelCase								=    '''
BERTScore Metrics with the hashcode from a source against one or more references.

Args:
    predictions (list of str): Prediction/candidate sentences.
    references (list of str or list of list of str): Reference sentences.
    lang (str): Language of the sentences; required (e.g. \'en\').
    model_type (str): Bert specification, default using the suggested
        model for the target language; has to specify at least one of
        `model_type` or `lang`.
    num_layers (int): The layer of representation to use,
        default using the number of layers tuned on WMT16 correlation data.
    verbose (bool): Turn on intermediate status update.
    idf (bool or dict): Use idf weighting; can also be a precomputed idf_dict.
    device (str): On which the contextual embedding model will be allocated on.
        If this argument is None, the model lives on cuda:0 if cuda is available.
    nthreads (int): Number of threads.
    batch_size (int): Bert score processing batch size,
        at least one of `model_type` or `lang`. `lang` needs to be
        specified when `rescale_with_baseline` is True.
    rescale_with_baseline (bool): Rescale bertscore with pre-computed baseline.
    baseline_path (str): Customized baseline file.
    use_fast_tokenizer (bool): `use_fast` parameter passed to HF tokenizer. New in version 0.3.10.

Returns:
    precision: Precision.
    recall: Recall.
    f1: F1 score.
    hashcode: Hashcode of the library.

Examples:

    >>> predictions = ["hello there", "general kenobi"]
    >>> references = ["hello there", "general kenobi"]
    >>> bertscore = datasets.load_metric("bertscore")
    >>> results = bertscore.compute(predictions=predictions, references=references, lang="en")
    >>> print([round(v, 2) for v in results["f1"]])
    [1.0, 1.0]
'''




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







       def 	a_		(				self)   ->   Dict:
            return datasets.MetricInfo(
                description=_DESCRIPTION,							citation=_CITATION,							homepage='https://github.com/Tiiiger/bert_score',							inputs_description=_KWARGS_DESCRIPTION,							features=datasets.Features(
                    {
                        'predictions': datasets.Value('string',							id='sequence'),
                        'references': datasets.Sequence(datasets.Value('string',							id='sequence'),							id='references'),
                    }),							codebase_urls=['https://github.com/Tiiiger/bert_score'],							reference_urls=[
                    'https://github.com/Tiiiger/bert_score',
                    'https://arxiv.org/abs/1904.09675',
                ],							)







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__=None,							lowerCAmelCase__=None,							lowerCAmelCase__=None,							lowerCAmelCase__=False,							lowerCAmelCase__=False,							lowerCAmelCase__=None,							lowerCAmelCase__=64,							lowerCAmelCase__=4,							lowerCAmelCase__=False,							lowerCAmelCase__=False,							lowerCAmelCase__=None,							lowerCAmelCase__=False,							)   ->   Union[str, Any]:
            snake_case_       =      bert_score.utils.get_hash
            snake_case_       =      bert_score.BERTScorer

            if version.parse(bert_score.__version__) >= version.parse('0.3.10'):
                 snake_case_       =      functools.partial(lowerCAmelCase__,							use_fast_tokenizer=lowerCAmelCase__)
                 snake_case_       =      functools.partial(lowerCAmelCase__,							use_fast_tokenizer=lowerCAmelCase__)
            elif use_fast_tokenizer:
                 raise ImportWarning(
                     'To use a fast tokenizer, the module `bert-score>=0.3.10` is required, and the current version of `bert-score` doesn\'t match this condition.\n'
                     'You can install it with `pip install "bert-score>=0.3.10"`.')

            if model_type is None:
                 assert lang is not None, "either lang or model_type should be specified"
                 snake_case_       =      bert_score.utils.langamodel[lang.lower()]

            if num_layers is None:
                 snake_case_       =      bert_score.utils.modelalayers[model_type]

            snake_case_       =      get_hash(
                model=lowerCAmelCase__,							num_layers=lowerCAmelCase__,							idf=lowerCAmelCase__,							rescale_with_baseline=lowerCAmelCase__,							use_custom_baseline=baseline_path is not None,							)

            with filter_logging_context():
                 if not hasattr(self,							'cached_bertscorer') or self.cached_bertscorer.hash != hashcode:
                      snake_case_       =      scorer(
                          model_type=lowerCAmelCase__,							num_layers=lowerCAmelCase__,							batch_size=lowerCAmelCase__,							nthreads=lowerCAmelCase__,							all_layers=lowerCAmelCase__,							idf=lowerCAmelCase__,							device=lowerCAmelCase__,							lang=lowerCAmelCase__,							rescale_with_baseline=lowerCAmelCase__,							baseline_path=lowerCAmelCase__,							)

            ((snake_case_) ,			(snake_case_) ,			(snake_case_))       =      self.cached_bertscorer.score(
                cands=lowerCAmelCase__,							refs=lowerCAmelCase__,							verbose=lowerCAmelCase__,							batch_size=lowerCAmelCase__,							)
            snake_case_       =      {
                'precision': P.tolist(),
                'recall': R.tolist(),
                'f1': F.tolist(),
                'hashcode': hashcode,
            }
            return output_dict







       def 	a_		(				self,							lowerCAmelCase__=None,							lowerCAmelCase__=None,							**lowerCAmelCase__)   ->   List[str]:
            # References can be strings or lists of strings
            # Let's change strings to lists of strings with one element
            if references is not None:
                 snake_case_       =      [[ref] if isinstance(lowerCAmelCase__,							lowerCAmelCase__) else ref for ref in references]
            super().add_batch(predictions=lowerCAmelCase__,							references=lowerCAmelCase__,							**lowerCAmelCase__)







       def 	a_		(				self,							lowerCAmelCase__=None,							lowerCAmelCase__=None,							**lowerCAmelCase__)   ->   List[str]:
            # References can be strings or lists of strings
            # Let's change strings to lists of strings with one element
            if isinstance(lowerCAmelCase__,							lowerCAmelCase__):
                 snake_case_       =      [reference]
            super().add(prediction=lowerCAmelCase__,							reference=lowerCAmelCase__,							**lowerCAmelCase__)


"""simple docstring"""


from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import Dataset

from .deepspeed import is_deepspeed_zeroa_enabled
from .generation.configuration_utils import GenerationConfig
from .trainer import Trainer
from .utils import logging


if TYPE_CHECKING:
    from .data.data_collator import DataCollator
    from .modeling_utils import PreTrainedModel
    from .tokenization_utils_base import PreTrainedTokenizerBase
    from .trainer_callback import TrainerCallback
    from .trainer_utils import EvalPrediction, PredictionOutput
    from .training_args import TrainingArguments


__UpperCamelCase								=    logging.get_logger(__name__)




class 				UpperCamelCase      (       lowerCAmelCase__     ):







       def __init__(				self,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = (None, None),							lowerCAmelCase__ = None,							)   ->   Dict:
            super().__init__(
                model=lowerCAmelCase__,							args=lowerCAmelCase__,							data_collator=lowerCAmelCase__,							train_dataset=lowerCAmelCase__,							eval_dataset=lowerCAmelCase__,							tokenizer=lowerCAmelCase__,							model_init=lowerCAmelCase__,							compute_metrics=lowerCAmelCase__,							callbacks=lowerCAmelCase__,							optimizers=lowerCAmelCase__,							preprocess_logits_for_metrics=lowerCAmelCase__,							)

            # Override self.model.generation_config if a GenerationConfig is specified in args.
            # Priority: args.generation_config > model.generation_config > default GenerationConfig.
            if self.args.generation_config is not None:
                 snake_case_       =      self.load_generation_config(self.args.generation_config)
                 snake_case_       =      gen_config







       @staticmethod
       def 	a_		(				lowerCAmelCase__)   ->   GenerationConfig:

            # GenerationConfig provided, nothing to do
            if isinstance(lowerCAmelCase__,							lowerCAmelCase__):
                 return deepcopy(lowerCAmelCase__)

            # str or Path
            snake_case_       =      Path(lowerCAmelCase__) if isinstance(lowerCAmelCase__,							lowerCAmelCase__) else gen_config_arg
            snake_case_       =      None

            # Figuring if it is path pointing to a file, pointing to a directory or else a model id or URL
            # This step is required in order to determine config_file_name
            if pretrained_model_name.is_file():
                 snake_case_       =      pretrained_model_name.name
                 snake_case_       =      pretrained_model_name.parent
            # dir path
            elif pretrained_model_name.is_dir():
                 pass
            # model id or URL
            else:
                 snake_case_       =      gen_config_arg

            snake_case_       =      GenerationConfig.from_pretrained(lowerCAmelCase__,							lowerCAmelCase__)
            return gen_config







       def 	a_		(				self,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = "eval",							**lowerCAmelCase__,							)   ->   Dict[str, float]:
            snake_case_       =      gen_kwargs.copy()
            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.args.generation_max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
            )
            snake_case_       =      gen_kwargs

            return super().evaluate(lowerCAmelCase__,							ignore_keys=lowerCAmelCase__,							metric_key_prefix=lowerCAmelCase__)







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__ = None,							lowerCAmelCase__ = "test",							**lowerCAmelCase__,							)   ->   "PredictionOutput":
            snake_case_       =      gen_kwargs.copy()
            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.args.generation_max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
            )
            snake_case_       =      gen_kwargs

            return super().predict(lowerCAmelCase__,							ignore_keys=lowerCAmelCase__,							metric_key_prefix=lowerCAmelCase__)







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

            if not self.args.predict_with_generate or prediction_loss_only:
                 return super().prediction_step(
                     lowerCAmelCase__,							lowerCAmelCase__,							prediction_loss_only=lowerCAmelCase__,							ignore_keys=lowerCAmelCase__)

            snake_case_       =      'labels' in inputs
            snake_case_       =      self._prepare_inputs(lowerCAmelCase__)

            # XXX: adapt synced_gpus for fairscale as well
            # Priority (handled in generate):
            # gen_kwargs > model.generation_config > default GenerationConfig()

            if len(lowerCAmelCase__) == 0 and hasattr(self,							'_gen_kwargs'):
                 snake_case_       =      self._gen_kwargs.copy()

            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.model.config.max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.model.config.num_beams
            )
            snake_case_       =      True if is_deepspeed_zeroa_enabled() else False
            snake_case_       =      (
                gen_kwargs['synced_gpus'] if gen_kwargs.get('synced_gpus') is not None else default_synced_gpus
            )

            # If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate
            # (otherwise, it would continue generating from the padded `decoder_input_ids`)
            if (
                "labels" in inputs
                and "decoder_input_ids" in inputs
                and inputs["labels"].shape == inputs["decoder_input_ids"].shape
            ):
                 snake_case_       =      {k: v for k, v in inputs.items() if k != 'decoder_input_ids'}
            snake_case_       =      self.model.generate(**lowerCAmelCase__,							**lowerCAmelCase__)

            # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
            # TODO: remove this hack when the legacy code that initializes generation_config from a model config is
            # removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183
            if self.model.generation_config._from_model_config:
                 snake_case_       =      False

            # Retrieves GenerationConfig from model.generation_config
            snake_case_       =      self.model.generation_config
            # in case the batch is shorter than max length, the output should be padded
            if generated_tokens.shape[-1] < gen_config.max_length:
                 snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_length)
            elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
                 snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_new_tokens + 1)

            with torch.no_grad():
                 if has_labels:
                      with self.compute_loss_context_manager():
                           snake_case_       =      model(**lowerCAmelCase__)
                      if self.label_smoother is not None:
                           snake_case_       =      self.label_smoother(lowerCAmelCase__,							inputs['labels']).mean().detach()
                      else:
                           snake_case_       =      (outputs['loss'] if isinstance(lowerCAmelCase__,							lowerCAmelCase__) else outputs[0]).mean().detach()
                 else:
                      snake_case_       =      None

            if self.args.prediction_loss_only:
                 return loss, None, None

            if has_labels:
                 snake_case_       =      inputs['labels']
                 if labels.shape[-1] < gen_config.max_length:
                      snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_length)
                 elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
                      snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_new_tokens + 1)
            else:
                 snake_case_       =      None

            return loss, generated_tokens, labels







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__)   ->   Any:
            if self.tokenizer is not None and hasattr(self.tokenizer,							'pad_token_id'):
                 # If PAD token is not defined at least EOS token has to be defined
                 snake_case_       =      (
                     self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
                 )
            else:
                 if self.model.config.pad_token_id is not None:
                      snake_case_       =      self.model.config.pad_token_id
                 else:
                      raise ValueError('Pad_token_id must be set in the configuration of the model, in order to pad tensors')

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

