



"""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 				UpperCamelCase__       (     lowercase_  ):




      """simple docstring"""




      def __init__(				self    :				List[Any]   ,      *SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				str=None   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]=None   ,      **SCREAMING_SNAKE_CASE_    :				List[str]	):
        super().__init__(*SCREAMING_SNAKE_CASE_   ,      **SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Optional[int]								=			eval_examples
        lowerCAmelCase_       :    List[str]								=			post_process_function



      def 					SCREAMING_SNAKE_CASE__			(				self    :				str   ,      SCREAMING_SNAKE_CASE_    :				Optional[Dataset] = None   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=None   ,      SCREAMING_SNAKE_CASE_    :				Optional[List[str]] = None   ,      SCREAMING_SNAKE_CASE_    :				str = "eval"   ,      **SCREAMING_SNAKE_CASE_    :				Tuple   ,      ):
        lowerCAmelCase_       :    Tuple								=			gen_kwargs.copy()
        lowerCAmelCase_       :    Union[str, Any]								=			(
            gen_kwargs['max_length'] if gen_kwargs.get('max_length'	) is not None else self.args.generation_max_length
        )
        lowerCAmelCase_       :    str								=			(
            gen_kwargs['num_beams'] if gen_kwargs.get('num_beams'	) is not None else self.args.generation_num_beams
        )
        lowerCAmelCase_       :    Optional[int]								=			gen_kwargs

        lowerCAmelCase_       :    int								=			self.eval_dataset if eval_dataset is None else eval_dataset
        lowerCAmelCase_       :    Optional[int]								=			self.get_eval_dataloader(SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    List[str]								=			self.eval_examples if eval_examples is None else eval_examples

        # Temporarily disable metric computation, we will do it in the loop here.
        lowerCAmelCase_       :    int								=			self.compute_metrics
        lowerCAmelCase_       :    List[Any]								=			None
        lowerCAmelCase_       :    List[str]								=			time.time()
        lowerCAmelCase_       :    Optional[Any]								=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
          lowerCAmelCase_       :    Dict								=			eval_loop(
              SCREAMING_SNAKE_CASE_   ,      description='Evaluation'   ,      prediction_loss_only=True if compute_metrics is None else None   ,      ignore_keys=SCREAMING_SNAKE_CASE_   ,      metric_key_prefix=SCREAMING_SNAKE_CASE_   ,      )
        finally:
          lowerCAmelCase_       :    str								=			compute_metrics
        lowerCAmelCase_       :    Dict								=			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(
                SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      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
          lowerCAmelCase_       :    List[str]								=			self.post_process_function(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    int								=			self.compute_metrics(SCREAMING_SNAKE_CASE_	)

          # Prefix all keys with metric_key_prefix + '_'
          for key in list(metrics.keys()	):
            if not key.startswith(F"{metric_key_prefix}_"	):
              lowerCAmelCase_       :    Union[str, Any]								=			metrics.pop(SCREAMING_SNAKE_CASE_	)

          metrics.update(output.metrics	)
        else:
          lowerCAmelCase_       :    Any								=			output.metrics

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

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

        lowerCAmelCase_       :    str								=			self.callback_handler.on_evaluate(self.args   ,      self.state   ,      self.control   ,      SCREAMING_SNAKE_CASE_	)
        return metrics





      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				str=None   ,      SCREAMING_SNAKE_CASE_    :				str = "test"   ,      **SCREAMING_SNAKE_CASE_    :				List[str]	):
        lowerCAmelCase_       :    Union[str, Any]								=			gen_kwargs.copy()

        lowerCAmelCase_       :    int								=			self.get_test_dataloader(SCREAMING_SNAKE_CASE_	)

        # Temporarily disable metric computation, we will do it in the loop here.
        lowerCAmelCase_       :    List[Any]								=			self.compute_metrics
        lowerCAmelCase_       :    Any								=			None
        lowerCAmelCase_       :    Any								=			time.time()
        lowerCAmelCase_       :    Union[str, Any]								=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
          lowerCAmelCase_       :    Optional[Any]								=			eval_loop(
              SCREAMING_SNAKE_CASE_   ,      description='Prediction'   ,      prediction_loss_only=True if compute_metrics is None else None   ,      ignore_keys=SCREAMING_SNAKE_CASE_   ,      metric_key_prefix=SCREAMING_SNAKE_CASE_   ,      )
        finally:
          lowerCAmelCase_       :    Dict								=			compute_metrics

        lowerCAmelCase_       :    Union[str, Any]								=			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(
                SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      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

        lowerCAmelCase_       :    int								=			self.post_process_function(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      'predict'	)
        lowerCAmelCase_       :    Any								=			self.compute_metrics(SCREAMING_SNAKE_CASE_	)

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










"""simple docstring"""







import logging

from transformers.configuration_utils import PretrainedConfig


lowercase__		:    List[Any]				   =     logging.getLogger(__name__)







class 				UpperCamelCase__       (     lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					"""masked_bert"""



      def __init__(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=3_0_5_2_2   ,      SCREAMING_SNAKE_CASE_    :				int=7_6_8   ,      SCREAMING_SNAKE_CASE_    :				Tuple=1_2   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=1_2   ,      SCREAMING_SNAKE_CASE_    :				Tuple=3_0_7_2   ,      SCREAMING_SNAKE_CASE_    :				List[str]="gelu"   ,      SCREAMING_SNAKE_CASE_    :				Tuple=0.1   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=5_1_2   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=2   ,      SCREAMING_SNAKE_CASE_    :				Dict=0.02   ,      SCREAMING_SNAKE_CASE_    :				Any=1E-12   ,      SCREAMING_SNAKE_CASE_    :				List[str]=0   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]="topK"   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]="constant"   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=0.0   ,      **SCREAMING_SNAKE_CASE_    :				Any   ,      ):
        super().__init__(pad_token_id=SCREAMING_SNAKE_CASE_   ,      **SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    List[Any]								=			vocab_size
        lowerCAmelCase_       :    str								=			hidden_size
        lowerCAmelCase_       :    Optional[int]								=			num_hidden_layers
        lowerCAmelCase_       :    Dict								=			num_attention_heads
        lowerCAmelCase_       :    List[str]								=			hidden_act
        lowerCAmelCase_       :    List[Any]								=			intermediate_size
        lowerCAmelCase_       :    Any								=			hidden_dropout_prob
        lowerCAmelCase_       :    str								=			attention_probs_dropout_prob
        lowerCAmelCase_       :    Any								=			max_position_embeddings
        lowerCAmelCase_       :    Dict								=			type_vocab_size
        lowerCAmelCase_       :    Tuple								=			initializer_range
        lowerCAmelCase_       :    List[Any]								=			layer_norm_eps
        lowerCAmelCase_       :    str								=			pruning_method
        lowerCAmelCase_       :    Optional[Any]								=			mask_init
        lowerCAmelCase_       :    int								=			mask_scale





