

import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    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.utils.versions import require_version


_SCREAMING_SNAKE_CASE  :						Any         =    logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

_SCREAMING_SNAKE_CASE  :						int         =    {
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
_SCREAMING_SNAKE_CASE  :						List[str]         =    {
    '''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,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
_SCREAMING_SNAKE_CASE  :						Tuple         =    sorted(arg_to_scheduler.keys())
_SCREAMING_SNAKE_CASE  :						int         =    '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''






class  UpperCAmelCase__	( pl.LightningModule ):






  """simple docstring"""


  def __init__(							self		:	Optional[int]			,     __lowerCamelCase		:	argparse.Namespace			,     __lowerCamelCase		:	List[str]=None			,     __lowerCamelCase		:	Any="base"			,     __lowerCamelCase		:	List[str]=None			,     __lowerCamelCase		:	Dict=None			,     __lowerCamelCase		:	Optional[int]=None			,     **__lowerCamelCase		:	Optional[Any]			,     )	->							Optional[int]:
       super().__init__()
       # TODO: move to self.save_hyperparameters()
       # self.save_hyperparameters()
       # can also expand arguments into trainer signature for easier reading

       self.save_hyperparameters(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					0
       SCREAMING_SNAKE_CASE__					=					Path(self.hparams.output_dir      )
       SCREAMING_SNAKE_CASE__					=					self.hparams.cache_dir if self.hparams.cache_dir else None
       if config is None:
            SCREAMING_SNAKE_CASE__					=					AutoConfig.from_pretrained(
                self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path			,     **({'''num_labels''': num_labels} if num_labels is not None else {})			,     cache_dir=__lowerCamelCase			,     **__lowerCamelCase			,     )
       else:
            SCREAMING_SNAKE_CASE__					=					config

       SCREAMING_SNAKE_CASE__					=					('''encoder_layerdrop''', '''decoder_layerdrop''', '''dropout''', '''attention_dropout''')
       for p in extra_model_params:
            if getattr(self.hparams			,     __lowerCamelCase			,     __lowerCamelCase      ):
                 assert hasattr(self.config			,     __lowerCamelCase      ), f'''model config doesn\'t have a `{p}` attribute'''
                 setattr(self.config			,     __lowerCamelCase			,     getattr(self.hparams			,     __lowerCamelCase      )      )

       if tokenizer is None:
            SCREAMING_SNAKE_CASE__					=					AutoTokenizer.from_pretrained(
                self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path			,     cache_dir=__lowerCamelCase			,     )
       else:
            SCREAMING_SNAKE_CASE__					=					tokenizer
       SCREAMING_SNAKE_CASE__					=					MODEL_MODES[mode]
       if model is None:
            SCREAMING_SNAKE_CASE__					=					self.model_type.from_pretrained(
                self.hparams.model_name_or_path			,     from_tf=bool('''.ckpt''' in self.hparams.model_name_or_path      )			,     config=self.config			,     cache_dir=__lowerCamelCase			,     )
       else:
            SCREAMING_SNAKE_CASE__					=					model






  def    lowercase_	(							self		:	Tuple			,     *__lowerCamelCase		:	List[Any]			,     **__lowerCamelCase		:	List[str]      )	->							Optional[int]:
       SCREAMING_SNAKE_CASE__					=					self.model_type.from_pretrained(*__lowerCamelCase			,     **__lowerCamelCase      )






  def    lowercase_	(							self		:	str      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					arg_to_scheduler[self.hparams.lr_scheduler]
       SCREAMING_SNAKE_CASE__					=					get_schedule_func(
           self.opt			,     num_warmup_steps=self.hparams.warmup_steps			,     num_training_steps=self.total_steps()      )
       SCREAMING_SNAKE_CASE__					=					{'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
       return scheduler






  def    lowercase_	(							self		:	str      )	->							List[Any]:
       SCREAMING_SNAKE_CASE__					=					self.model
       SCREAMING_SNAKE_CASE__					=					['''bias''', '''LayerNorm.weight''']
       SCREAMING_SNAKE_CASE__					=					[
           {
               '''params''': [
                   p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay      )
               ],  # check this named paramters
               '''weight_decay''': self.hparams.weight_decay,
           },
           {
               '''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay      )],
               '''weight_decay''': 0.0,
           },
       ]
       if self.hparams.adafactor:
            SCREAMING_SNAKE_CASE__					=					Adafactor(
                __lowerCamelCase			,     lr=self.hparams.learning_rate			,     scale_parameter=__lowerCamelCase			,     relative_step=__lowerCamelCase      )

       else:
            SCREAMING_SNAKE_CASE__					=					AdamW(
                __lowerCamelCase			,     lr=self.hparams.learning_rate			,     eps=self.hparams.adam_epsilon      )
       SCREAMING_SNAKE_CASE__					=					optimizer

       SCREAMING_SNAKE_CASE__					=					self.get_lr_scheduler()

       return [optimizer], [scheduler]






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	Tuple      )	->							Dict:
       return self.validation_step(__lowerCamelCase			,     __lowerCamelCase      )






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	List[Any]      )	->							Union[str, Any]:
       return self.validation_end(__lowerCamelCase      )






  def    lowercase_	(							self		:	List[Any]      )	->							int:
       SCREAMING_SNAKE_CASE__					=					max(1			,     self.hparams.gpus      )  # TODO: consider num_tpu_cores
       SCREAMING_SNAKE_CASE__					=					self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
       return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs






  def    lowercase_	(							self		:	str			,     __lowerCamelCase		:	Any      )	->							Dict:
       if stage == "test":
            SCREAMING_SNAKE_CASE__					=					len(self.test_dataloader().dataset      )
       else:
            SCREAMING_SNAKE_CASE__					=					self.get_dataloader('''train'''			,     self.hparams.train_batch_size			,     shuffle=__lowerCamelCase      )
            SCREAMING_SNAKE_CASE__					=					len(self.train_dataloader().dataset      )






  def    lowercase_	(							self		:	Optional[Any]			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	bool = False      )	->							Dict:
       raise NotImplementedError('''You must implement this for your task'''      )






  def    lowercase_	(							self		:	Any      )	->							int:
       return self.train_loader






  def    lowercase_	(							self		:	Optional[Any]      )	->							int:
       return self.get_dataloader('''dev'''			,     self.hparams.eval_batch_size			,     shuffle=__lowerCamelCase      )






  def    lowercase_	(							self		:	Dict      )	->							Optional[Any]:
       return self.get_dataloader('''test'''			,     self.hparams.eval_batch_size			,     shuffle=__lowerCamelCase      )






  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	List[Any]      )	->							Optional[Any]:
       return os.path.join(
           self.hparams.data_dir			,     '''cached_{}_{}_{}'''.format(
               __lowerCamelCase			,     list(filter(__lowerCamelCase			,     self.hparams.model_name_or_path.split('''/'''      )      )      ).pop()			,     str(self.hparams.max_seq_length      )			,     )			,     )






  @pl.utilities.rank_zero_only
  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Dict[str, Any]      )	->							None:
       SCREAMING_SNAKE_CASE__					=					self.output_dir.joinpath('''best_tfmr'''      )
       SCREAMING_SNAKE_CASE__					=					self.step_count
       self.model.save_pretrained(__lowerCamelCase      )
       self.tokenizer.save_pretrained(__lowerCamelCase      )







  @staticmethod
  def    lowercase_	(							__lowerCamelCase		:	Union[str, Any]			,     __lowerCamelCase		:	List[str]      )	->							Union[str, Any]:
       parser.add_argument(
           '''--model_name_or_path'''			,     default=__lowerCamelCase			,     type=__lowerCamelCase			,     required=__lowerCamelCase			,     help='''Path to pretrained model or model identifier from huggingface.co/models'''			,     )
       parser.add_argument(
           '''--config_name'''			,     default=''''''			,     type=__lowerCamelCase			,     help='''Pretrained config name or path if not the same as model_name'''      )
       parser.add_argument(
           '''--tokenizer_name'''			,     default=__lowerCamelCase			,     type=__lowerCamelCase			,     help='''Pretrained tokenizer name or path if not the same as model_name'''			,     )
       parser.add_argument(
           '''--cache_dir'''			,     default=str(Path(__lowerCamelCase      ).parent / '''test_run''' / '''cache'''      )			,     type=__lowerCamelCase			,     help='''Where do you want to store the pre-trained models downloaded from huggingface.co'''			,     )
       parser.add_argument(
           '''--encoder_layerdrop'''			,     type=__lowerCamelCase			,     help='''Encoder layer dropout probability (Optional). Goes into model.config'''			,     )
       parser.add_argument(
           '''--decoder_layerdrop'''			,     type=__lowerCamelCase			,     help='''Decoder layer dropout probability (Optional). Goes into model.config'''			,     )
       parser.add_argument(
           '''--dropout'''			,     type=__lowerCamelCase			,     help='''Dropout probability (Optional). Goes into model.config'''			,     )
       parser.add_argument(
           '''--attention_dropout'''			,     type=__lowerCamelCase			,     help='''Attention dropout probability (Optional). Goes into model.config'''			,     )
       parser.add_argument('''--learning_rate'''			,     default=5e-5			,     type=__lowerCamelCase			,     help='''The initial learning rate for Adam.'''      )
       parser.add_argument(
           '''--lr_scheduler'''			,     default='''linear'''			,     choices=__lowerCamelCase			,     metavar=__lowerCamelCase			,     type=__lowerCamelCase			,     help='''Learning rate scheduler'''			,     )
       parser.add_argument('''--weight_decay'''			,     default=0.0			,     type=__lowerCamelCase			,     help='''Weight decay if we apply some.'''      )
       parser.add_argument('''--adam_epsilon'''			,     default=1e-8			,     type=__lowerCamelCase			,     help='''Epsilon for Adam optimizer.'''      )
       parser.add_argument('''--warmup_steps'''			,     default=0			,     type=__lowerCamelCase			,     help='''Linear warmup over warmup_steps.'''      )
       parser.add_argument('''--num_workers'''			,     default=4			,     type=__lowerCamelCase			,     help='''kwarg passed to DataLoader'''      )
       parser.add_argument('''--num_train_epochs'''			,     dest='''max_epochs'''			,     default=3			,     type=__lowerCamelCase      )
       parser.add_argument('''--train_batch_size'''			,     default=32			,     type=__lowerCamelCase      )
       parser.add_argument('''--eval_batch_size'''			,     default=32			,     type=__lowerCamelCase      )
       parser.add_argument('''--adafactor'''			,     action='''store_true'''      )






class  UpperCAmelCase__	( pl.Callback ):






  """simple docstring"""


  def    lowercase_	(							self		:	Union[str, Any]			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	Any      )	->							List[str]:
       if (
           trainer.is_global_zero and trainer.global_rank == 0
       ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
            pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.






class  UpperCAmelCase__	( pl.Callback ):






  """simple docstring"""


  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	str      )	->							Dict:
       # print(pl_module.model.rag)
       for name, param in pl_module.model.rag.named_parameters():
            if param.grad is None:
                 print(__lowerCamelCase      )






class  UpperCAmelCase__	( pl.Callback ):






  """simple docstring"""


  def    lowercase_	(							self		:	List[str]			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	Optional[int]      )	->							Dict:
       SCREAMING_SNAKE_CASE__					=					trainer.lr_schedulers[0]['''scheduler''']
       SCREAMING_SNAKE_CASE__					=					{f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()      )}
       pl_module.logger.log_metrics(__lowerCamelCase      )






  def    lowercase_	(							self		:	List[Any]			,     __lowerCamelCase		:	pl.Trainer			,     __lowerCamelCase		:	pl.LightningModule      )	->							Tuple:
       rank_zero_info('''***** Validation results *****'''      )
       SCREAMING_SNAKE_CASE__					=					trainer.callback_metrics
       # Log results
       for key in sorted(__lowerCamelCase      ):
            if key not in ["log", "progress_bar"]:
                 rank_zero_info('''{} = {}\n'''.format(__lowerCamelCase			,     str(metrics[key]      )      )      )







  def    lowercase_	(							self		:	List[Any]			,     __lowerCamelCase		:	pl.Trainer			,     __lowerCamelCase		:	pl.LightningModule      )	->							List[str]:
       rank_zero_info('''***** Test results *****'''      )
       SCREAMING_SNAKE_CASE__					=					trainer.callback_metrics
       # Log and save results to file
       SCREAMING_SNAKE_CASE__					=					os.path.join(pl_module.hparams.output_dir			,     '''test_results.txt'''      )
       with open(__lowerCamelCase			,     '''w'''      ) as writer:
            for key in sorted(__lowerCamelCase      ):
                 if key not in ["log", "progress_bar"]:
                      rank_zero_info('''{} = {}\n'''.format(__lowerCamelCase			,     str(metrics[key]      )      )      )
                      writer.write('''{} = {}\n'''.format(__lowerCamelCase			,     str(metrics[key]      )      )      )



def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''simple docstring'''




     parser.add_argument(
         '''--output_dir'''   ,	default=str(Path(_A       ).parent / '''test_run''' / '''model_checkpoints'''       )   ,	type=_A   ,	help='''The output directory where the model predictions and checkpoints will be written.'''   ,	)
     parser.add_argument(
         '''--fp16'''   ,	action='''store_true'''   ,	help='''Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'''   ,	)

     parser.add_argument(
         '''--fp16_opt_level'''   ,	type=_A   ,	default='''O2'''   ,	help=(
             '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'''
             '''See details at https://nvidia.github.io/apex/amp.html'''
         )   ,	)
     parser.add_argument('''--n_tpu_cores'''   ,	dest='''tpu_cores'''   ,	type=_A       )
     parser.add_argument('''--max_grad_norm'''   ,	dest='''gradient_clip_val'''   ,	default=1.0   ,	type=_A   ,	help='''Max gradient norm'''       )
     parser.add_argument('''--do_train'''   ,	action='''store_true'''   ,	help='''Whether to run training.'''       )
     parser.add_argument('''--do_predict'''   ,	action='''store_true'''   ,	help='''Whether to run predictions on the test set.'''       )
     parser.add_argument(
         '''--gradient_accumulation_steps'''   ,	dest='''accumulate_grad_batches'''   ,	type=_A   ,	default=1   ,	help='''Number of updates steps to accumulate before performing a backward/update pass.'''   ,	)
     parser.add_argument('''--seed'''   ,	type=_A   ,	default=42   ,	help='''random seed for initialization'''       )
     parser.add_argument(
         '''--data_dir'''   ,	default=str(Path(_A       ).parent / '''test_run''' / '''dummy-train-data'''       )   ,	type=_A   ,	help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.'''   ,	)







def        UpperCAmelCase_     (		_A   ,	_A   ,	_A=None   ,	_A=True   ,	_A=[]   ,	_A=None   ,	_A=None   ,	**_A   ,	):





     '''simple docstring'''




     pl.seed_everything(args.seed       )

     # init model
     SCREAMING_SNAKE_CASE__					=					Path(model.hparams.output_dir       )
     odir.mkdir(exist_ok=_A       )

     # add custom checkpoints
     if checkpoint_callback is None:
          SCREAMING_SNAKE_CASE__					=					pl.callbacks.ModelCheckpoint(
              filepath=args.output_dir   ,	prefix='''checkpoint'''   ,	monitor='''val_loss'''   ,	mode='''min'''   ,	save_top_k=1       )
     if early_stopping_callback:
          extra_callbacks.append(_A       )
     if logging_callback is None:
          SCREAMING_SNAKE_CASE__					=					LoggingCallback()

     SCREAMING_SNAKE_CASE__					=					{}

     if args.fpaa:
          SCREAMING_SNAKE_CASE__					=					16

     if args.gpus > 1:
          SCREAMING_SNAKE_CASE__					=					'''auto'''
          SCREAMING_SNAKE_CASE__					=					'''ddp'''

     SCREAMING_SNAKE_CASE__					=					args.accumulate_grad_batches
     SCREAMING_SNAKE_CASE__					=					None
     SCREAMING_SNAKE_CASE__					=					'''auto'''

     SCREAMING_SNAKE_CASE__					=					pl.Trainer.from_argparse_args(
         _A   ,	weights_summary=_A   ,	callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]   ,	logger=_A   ,	val_check_interval=1   ,	num_sanity_val_steps=2   ,	**_A   ,	)

     if args.do_train:
          trainer.fit(_A       )

     else:
          print('''RAG modeling tests with new set functions successfuly executed!'''       )
     return trainer





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  UpperCAmelCase__	( A__ ):






  """simple docstring"""


  def __init__(							self		:	Tuple			,     *__lowerCamelCase		:	List[str]			,     __lowerCamelCase		:	Any=None			,     __lowerCamelCase		:	Optional[int]=None			,     **__lowerCamelCase		:	Optional[int]      )	->							List[str]:
       super().__init__(*__lowerCamelCase			,     **__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					eval_examples
       SCREAMING_SNAKE_CASE__					=					post_process_function






  def    lowercase_	(							self		:	Optional[Any]			,     __lowerCamelCase		:	Dict=None			,     __lowerCamelCase		:	int=None			,     __lowerCamelCase		:	str=None			,     __lowerCamelCase		:	str = "eval"      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					self.eval_dataset if eval_dataset is None else eval_dataset
       SCREAMING_SNAKE_CASE__					=					self.get_eval_dataloader(__lowerCamelCase      )
       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(
                __lowerCamelCase			,     description='''Evaluation'''			,     prediction_loss_only=True if compute_metrics is None else None			,     ignore_keys=__lowerCamelCase			,     metric_key_prefix=__lowerCamelCase			,     )
       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(
               __lowerCamelCase			,     __lowerCamelCase			,     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(__lowerCamelCase			,     __lowerCamelCase			,     output.predictions      )
            SCREAMING_SNAKE_CASE__					=					self.compute_metrics(__lowerCamelCase      )

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

       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			,     __lowerCamelCase      )
       return metrics







  def    lowercase_	(							self		:	Dict			,     __lowerCamelCase		:	Dict			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	str=None			,     __lowerCamelCase		:	str = "test"      )	->							Any:
       SCREAMING_SNAKE_CASE__					=					self.get_test_dataloader(__lowerCamelCase      )

       # 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(
                __lowerCamelCase			,     description='''Prediction'''			,     prediction_loss_only=True if compute_metrics is None else None			,     ignore_keys=__lowerCamelCase			,     metric_key_prefix=__lowerCamelCase			,     )
       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(
               __lowerCamelCase			,     __lowerCamelCase			,     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(__lowerCamelCase			,     __lowerCamelCase			,     output.predictions			,     '''predict'''      )
       SCREAMING_SNAKE_CASE__					=					self.compute_metrics(__lowerCamelCase      )

       # 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(__lowerCamelCase      )
       metrics.update(output.metrics      )
       return PredictionOutput(predictions=predictions.predictions			,     label_ids=predictions.label_ids			,     metrics=__lowerCamelCase      )


