





"""simple docstring"""





import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments

import datasets


lowerCAmelCase__        =      '''\
@article{eddine2021frugalscore,
  title={FrugalScore: Learning Cheaper, Lighter and Faster Evaluation Metrics for Automatic Text Generation},
  author={Eddine, Moussa Kamal and Shang, Guokan and Tixier, Antoine J-P and Vazirgiannis, Michalis},
  journal={arXiv preprint arXiv:2110.08559},
  year={2021}
}
'''

lowerCAmelCase__        =      '''\
FrugalScore is a reference-based metric for NLG models evaluation. It is based on a distillation approach that allows to learn a fixed, low cost version of any expensive NLG metric, while retaining most of its original performance.
'''


lowerCAmelCase__        =      '''
Calculates how good are predictions given some references, using certain scores.
Args:
    predictions (list of str): list of predictions to score. Each predictions
        should be a string.
    references (list of str): list of reference for each prediction. Each
        reference should be a string.
    batch_size (int): the batch size for predictions.
    max_length (int): maximum sequence length.
    device (str): either gpu or cpu
Returns:
    scores (list of int): list of scores.
Examples:
    >>> frugalscore = datasets.load_metric("frugalscore")
    >>> results = frugalscore.compute(predictions=[\'hello there\', \'huggingface\'], references=[\'hello world\', \'hugging face\'])
    >>> print([round(s, 3) for s in results["scores"]])
    [0.631, 0.645]
'''





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


						def      snake_case_ (self		)			-> Tuple:
								return datasets.MetricInfo(
								    description=_DESCRIPTION ,				citation=_CITATION ,				inputs_description=_KWARGS_DESCRIPTION ,				features=datasets.Features(
								        {
								            "predictions": datasets.Value("string"		),
								            "references": datasets.Value("string"		),
								        }		) ,				homepage="https://github.com/moussaKam/FrugalScore" ,				)


						def      snake_case_ (self ,				__a		)			-> Optional[int]:
								if self.config_name == "default":
										UpperCamelCase =						"moussaKam/frugalscore_tiny_bert-base_bert-score"
								else:
										UpperCamelCase =						self.config_name
								UpperCamelCase =						AutoModelForSequenceClassification.from_pretrained(__a		)
								UpperCamelCase =						AutoTokenizer.from_pretrained(__a		)




						def      snake_case_ (self ,				__a ,				__a ,				__a=32 ,				__a=1_28 ,				__a=None ,				)			-> Optional[int]:
								assert len(__a		) == len(
								    __a		), "predictions and references should have the same number of sentences."
								if device is not None:
										assert device in ["gpu", "cpu"], "device should be either gpu or cpu."
								else:
										UpperCamelCase =						"gpu" if torch.cuda.is_available() else "cpu"
								UpperCamelCase =						TrainingArguments(
								    "trainer" ,				fpaa=(device == "gpu") ,				per_device_eval_batch_size=__a ,				report_to="all" ,				no_cuda=(device == "cpu") ,				log_level="warning" ,				)
								UpperCamelCase =						{"sentence1": predictions, "sentence2": references}
								UpperCamelCase =						datasets.Dataset.from_dict(__a		)

								def tokenize_function(__a		):
										return self.tokenizer(
										    data["sentence1"] ,				data["sentence2"] ,				max_length=__a ,				truncation=__a ,				padding=__a		)

								UpperCamelCase =						raw_datasets.map(__a ,				batched=__a		)
								tokenized_datasets.remove_columns(["sentence1", "sentence2"]		)
								UpperCamelCase =						Trainer(self.model ,				__a ,				tokenizer=self.tokenizer		)
								UpperCamelCase =						trainer.predict(__a		)
								return {"scores": list(predictions.predictions.squeeze(-1		)		)}













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


lowerCAmelCase__        =      logging.get_logger(__name__)





class 							_lowerCamelCase   (							_lowercase			):


						def __init__(self ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = (None, None) ,				__a = None ,				)			-> List[Any]:
								super().__init__(
								    model=__a ,				args=__a ,				data_collator=__a ,				train_dataset=__a ,				eval_dataset=__a ,				tokenizer=__a ,				model_init=__a ,				compute_metrics=__a ,				callbacks=__a ,				optimizers=__a ,				preprocess_logits_for_metrics=__a ,				)

								# 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:
										UpperCamelCase =						self.load_generation_config(self.args.generation_config		)
										UpperCamelCase =						gen_config


						@staticmethod
						def      snake_case_ (__a		)			-> GenerationConfig:

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

								# str or Path
								UpperCamelCase =						Path(__a		) if isinstance(__a ,				__a		) else gen_config_arg
								UpperCamelCase =						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():
										UpperCamelCase =						pretrained_model_name.name
										UpperCamelCase =						pretrained_model_name.parent
								# dir path
								elif pretrained_model_name.is_dir():
										pass
								# model id or URL
								else:
										UpperCamelCase =						gen_config_arg

								UpperCamelCase =						GenerationConfig.from_pretrained(__a ,				__a		)
								return gen_config


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

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


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

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


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

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

								UpperCamelCase =						"labels" in inputs
								UpperCamelCase =						self._prepare_inputs(__a		)

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

								if len(__a		) == 0 and hasattr(self ,				"_gen_kwargs"		):
										UpperCamelCase =						self._gen_kwargs.copy()

								if gen_kwargs.get("max_length"		) is None and gen_kwargs.get("max_new_tokens"		) is None:
										UpperCamelCase =						self.model.config.max_length
								UpperCamelCase =						(
								    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"		) is not None else self.model.config.num_beams
								)
								UpperCamelCase =						True if is_deepspeed_zeroa_enabled() else False
								UpperCamelCase =						(
								    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
								):
										UpperCamelCase =						{k: v for k, v in inputs.items() if k != "decoder_input_ids"}
								UpperCamelCase =						self.model.generate(**__a ,				**__a		)

								# 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:
										UpperCamelCase =						False

								# Retrieves GenerationConfig from model.generation_config
								UpperCamelCase =						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:
										UpperCamelCase =						self._pad_tensors_to_max_len(__a ,				gen_config.max_length		)
								elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
										UpperCamelCase =						self._pad_tensors_to_max_len(__a ,				gen_config.max_new_tokens + 1		)

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

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

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

								return loss, generated_tokens, labels




						def      snake_case_ (self ,				__a ,				__a		)			-> str:
								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
										UpperCamelCase =						(
										    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:
												UpperCamelCase =						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"		)

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






