





"""simple docstring"""





import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm
						import torch_xla.debug.metrics as met





class        _lowerCamelCase      (       _lowercase ):


						def __init__(self   ,							*__a   ,							__a=None   ,							__a=None   ,							**__a		) ->   Tuple:
								super().__init__(*__a   ,							**__a		)
								UpperCamelCase   =	eval_examples
								UpperCamelCase   =	post_process_function


						def     snake_case_   (self   ,							__a = None   ,							__a=None   ,							__a = None   ,							__a = "eval"   ,							**__a   ,							) ->   Dict[str, float]:
								UpperCamelCase   =	gen_kwargs.copy()
								UpperCamelCase   =	(
								    gen_kwargs["max_length"] if gen_kwargs.get("max_length"		) is not None else 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

								UpperCamelCase   =	self.eval_dataset if eval_dataset is None else eval_dataset
								UpperCamelCase   =	self.get_eval_dataloader(__a		)
								UpperCamelCase   =	self.eval_examples if eval_examples is None else eval_examples

								# Temporarily disable metric computation, we will do it in the loop here.
								UpperCamelCase   =	self.compute_metrics
								UpperCamelCase   =	None
								UpperCamelCase   =	time.time()
								UpperCamelCase   =	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								try:
										UpperCamelCase   =	eval_loop(
										    __a   ,							description="Evaluation"   ,							prediction_loss_only=True if compute_metrics is None else None   ,							ignore_keys=__a   ,							metric_key_prefix=__a   ,							)
								finally:
										UpperCamelCase   =	compute_metrics
								UpperCamelCase   =	self.args.eval_batch_size * self.args.world_size
								if F"{metric_key_prefix}_jit_compilation_time" in output.metrics:
										start_time += output.metrics[F"{metric_key_prefix}_jit_compilation_time"]
								output.metrics.update(
								    speed_metrics(
								        __a   ,							__a   ,							num_samples=output.num_samples   ,							num_steps=math.ceil(output.num_samples / total_batch_size		)   ,							)		)

								if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
										# Only the main node write the results by default
										UpperCamelCase   =	self.post_process_function(__a   ,							__a   ,							__a		)
										UpperCamelCase   =	self.compute_metrics(__a		)

										# Prefix all keys with metric_key_prefix + '_'
										for key in list(metrics.keys()		):
												if not key.startswith(F"{metric_key_prefix}_"		):
														UpperCamelCase   =	metrics.pop(__a		)

										metrics.update(output.metrics		)
								else:
										UpperCamelCase   =	output.metrics

								if self.args.should_log:
										# Only the main node log the results by default
										self.log(__a		)

								if self.args.tpu_metrics_debug or self.args.debug:
										# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
										xm.master_print(met.metrics_report()		)

								UpperCamelCase   =	self.callback_handler.on_evaluate(self.args   ,							self.state   ,							self.control   ,							__a		)
								return metrics




						def     snake_case_   (self   ,							__a   ,							__a   ,							__a=None   ,							__a = "test"   ,							**__a		) ->   str:
								UpperCamelCase   =	gen_kwargs.copy()

								UpperCamelCase   =	self.get_test_dataloader(__a		)

								# Temporarily disable metric computation, we will do it in the loop here.
								UpperCamelCase   =	self.compute_metrics
								UpperCamelCase   =	None
								UpperCamelCase   =	time.time()
								UpperCamelCase   =	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								try:
										UpperCamelCase   =	eval_loop(
										    __a   ,							description="Prediction"   ,							prediction_loss_only=True if compute_metrics is None else None   ,							ignore_keys=__a   ,							metric_key_prefix=__a   ,							)
								finally:
										UpperCamelCase   =	compute_metrics

								UpperCamelCase   =	self.args.eval_batch_size * self.args.world_size
								if F"{metric_key_prefix}_jit_compilation_time" in output.metrics:
										start_time += output.metrics[F"{metric_key_prefix}_jit_compilation_time"]
								output.metrics.update(
								    speed_metrics(
								        __a   ,							__a   ,							num_samples=output.num_samples   ,							num_steps=math.ceil(output.num_samples / total_batch_size		)   ,							)		)
								if self.post_process_function is None or self.compute_metrics is None:
										return output

								UpperCamelCase   =	self.post_process_function(__a   ,							__a   ,							__a   ,							"predict"		)
								UpperCamelCase   =	self.compute_metrics(__a		)

								# Prefix all keys with metric_key_prefix + '_'
								for key in list(metrics.keys()		):
										if not key.startswith(F"{metric_key_prefix}_"		):
												UpperCamelCase   =	metrics.pop(__a		)
								metrics.update(output.metrics		)
								return PredictionOutput(predictions=predictions.predictions   ,							label_ids=predictions.label_ids   ,							metrics=__a		)













"""simple docstring"""





import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
						from transformers import AutoModelForSeqaSeqLM, AutoTokenizer





@require_torch
@require_sentencepiece
@require_tokenizers
class        _lowerCamelCase      (       unittest.TestCase ):


						@slow
						def     snake_case_   (self		) ->   List[str]:
								UpperCamelCase   =	AutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"   ,							return_dict=__a		).to(__a		)
								UpperCamelCase   =	AutoTokenizer.from_pretrained("google/mt5-small"		)

								UpperCamelCase   =	tokenizer("Hello there"   ,							return_tensors="pt"		).input_ids
								UpperCamelCase   =	tokenizer("Hi I am"   ,							return_tensors="pt"		).input_ids

								UpperCamelCase   =	model(input_ids.to(__a		)   ,							labels=labels.to(__a		)		).loss
								UpperCamelCase   =	-(labels.shape[-1] * loss.item())

								UpperCamelCase   =	-84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE		) < 1e-4		)






