import math
import time

from transformers import Trainer, 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 __snake_case				(						lowerCamelCase_		):




   def __init__(    self			:				Tuple       ,	*_lowercase			:				Dict       ,	_lowercase			:				Any=None       ,	_lowercase			:				Any=None       ,	**_lowercase			:				Any      ):



     """simple docstring"""






     super().__init__(*_lowercase       ,	**_lowercase      )
     SCREAMING_SNAKE_CASE__										=      eval_examples
     SCREAMING_SNAKE_CASE__										=      post_process_function




   def 	__a						(    self			:				str       ,	_lowercase			:				List[str]=None       ,	_lowercase			:				Union[str, Any]=None       ,	_lowercase			:				Optional[int]=None       ,	_lowercase			:				str = "eval"      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.eval_dataset if eval_dataset is None else eval_dataset
     SCREAMING_SNAKE_CASE__										=      self.get_eval_dataloader(_lowercase      )
     SCREAMING_SNAKE_CASE__										=      self.eval_examples if eval_examples is None else eval_examples

     # Temporarily disable metric computation, we will do it in the loop here.
     SCREAMING_SNAKE_CASE__										=      self.compute_metrics
     SCREAMING_SNAKE_CASE__										=      None
     SCREAMING_SNAKE_CASE__										=      self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
     SCREAMING_SNAKE_CASE__										=      time.time()
     try:
       SCREAMING_SNAKE_CASE__										=      eval_loop(
           _lowercase       ,	description="""Evaluation"""       ,	prediction_loss_only=True if compute_metrics is None else None       ,	ignore_keys=_lowercase       ,	metric_key_prefix=_lowercase       ,	)
     finally:
       SCREAMING_SNAKE_CASE__										=      compute_metrics
     SCREAMING_SNAKE_CASE__										=      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(
             _lowercase       ,	_lowercase       ,	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
       SCREAMING_SNAKE_CASE__										=      self.post_process_function(_lowercase       ,	_lowercase       ,	output.predictions      )
       SCREAMING_SNAKE_CASE__										=      self.compute_metrics(_lowercase      )

       # Prefix all keys with metric_key_prefix + '_'
       for key in list(metrics.keys()      ):
         if not key.startswith(f"""{metric_key_prefix}_"""      ):
           SCREAMING_SNAKE_CASE__										=      metrics.pop(_lowercase      )
       metrics.update(output.metrics      )
     else:
       SCREAMING_SNAKE_CASE__										=      output.metrics

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

     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()      )

     SCREAMING_SNAKE_CASE__										=      self.callback_handler.on_evaluate(self.args       ,	self.state       ,	self.control       ,	_lowercase      )
     return metrics




   def 	__a						(    self			:				int       ,	_lowercase			:				Optional[Any]       ,	_lowercase			:				str       ,	_lowercase			:				Any=None       ,	_lowercase			:				str = "test"      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      self.get_test_dataloader(_lowercase      )

     # Temporarily disable metric computation, we will do it in the loop here.
     SCREAMING_SNAKE_CASE__										=      self.compute_metrics
     SCREAMING_SNAKE_CASE__										=      None
     SCREAMING_SNAKE_CASE__										=      self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
     SCREAMING_SNAKE_CASE__										=      time.time()
     try:
       SCREAMING_SNAKE_CASE__										=      eval_loop(
           _lowercase       ,	description="""Prediction"""       ,	prediction_loss_only=True if compute_metrics is None else None       ,	ignore_keys=_lowercase       ,	metric_key_prefix=_lowercase       ,	)
     finally:
       SCREAMING_SNAKE_CASE__										=      compute_metrics
     SCREAMING_SNAKE_CASE__										=      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(
             _lowercase       ,	_lowercase       ,	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

     SCREAMING_SNAKE_CASE__										=      self.post_process_function(_lowercase       ,	_lowercase       ,	output.predictions       ,	"""predict"""      )
     SCREAMING_SNAKE_CASE__										=      self.compute_metrics(_lowercase      )

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

import warnings

from ..trainer import Trainer
from ..utils import logging


__lowerCamelCase			: List[Any]							   =		logging.get_logger(__name__)







class __snake_case				(						lowerCamelCase_		):




   def __init__(    self			:				Tuple       ,	_lowercase			:				Optional[int]=None       ,	**_lowercase			:				List[Any]      ):



     """simple docstring"""






     warnings.warn(
         """`SageMakerTrainer` is deprecated and will be removed in v5 of Transformers. You can use `Trainer` """
         """instead."""       ,	_lowercase       ,	)
     super().__init__(args=_lowercase       ,	**_lowercase      )
