


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 lowerCAmelCase_		(				UpperCamelCase_ ):

	'''simple docstring'''

	def __init__(			self	,   *_UpperCAmelCase	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   **_UpperCAmelCase      ):
							super().__init__(*_a	,   **_a      )
							snake_case_ =   eval_examples
							snake_case_ =   post_process_function



	def 		UpperCamelCase__  (			self	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   _UpperCAmelCase = "eval"      ):
							snake_case_ =   self.eval_dataset if eval_dataset is None else eval_dataset
							snake_case_ =   self.get_eval_dataloader(_a      )
							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.
							snake_case_ =   self.compute_metrics
							snake_case_ =   None
							snake_case_ =   self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
							snake_case_ =   time.time()
							try:
													snake_case_ =   eval_loop(
													    _a	,   description='''Evaluation'''	,   prediction_loss_only=True if compute_metrics is None else None	,   ignore_keys=_a	,   metric_key_prefix=_a	,   )
							finally:
													snake_case_ =   compute_metrics
							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(
							        _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
													snake_case_ =   self.post_process_function(_a	,   _a	,   output.predictions      )
													snake_case_ =   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}_'''      ):
																									snake_case_ =   metrics.pop(_a      )
													metrics.update(output.metrics      )
							else:
													snake_case_ =   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()      )

							snake_case_ =   self.callback_handler.on_evaluate(self.args	,   self.state	,   self.control	,   _a      )
							return metrics





	def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase=None	,   _UpperCAmelCase = "test"      ):
							snake_case_ =   self.get_test_dataloader(_a      )

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

							snake_case_ =   self.post_process_function(_a	,   _a	,   output.predictions	,   '''predict'''      )
							snake_case_ =   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}_'''      ):
																			snake_case_ =   metrics.pop(_a      )
							metrics.update(output.metrics      )
							return PredictionOutput(predictions=predictions.predictions	,   label_ids=predictions.label_ids	,   metrics=_a      )



from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput







@dataclass
class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						None
 __snake_case				=						None
 __snake_case				=						None
 __snake_case				=						None







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase=1	,   _UpperCAmelCase=0	,   _UpperCAmelCase=2	,   _UpperCAmelCase=5_12	,   _UpperCAmelCase="cls"	,   _UpperCAmelCase=False	,   _UpperCAmelCase=True	,   **_UpperCAmelCase	,   ):
       super().__init__(pad_token_id=_UpperCAmelCase	,   bos_token_id=_UpperCAmelCase	,   eos_token_id=_UpperCAmelCase	,   **_UpperCAmelCase      )
       snake_case_ =   project_dim
       snake_case_ =   pooler_fn
       snake_case_ =   learn_encoder
       snake_case_ =   use_attention_mask







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						[R"pooler", R"logit_scale"]
 __snake_case				=						[R"position_ids", R"predictions.decoder.bias"]
 __snake_case				=						"roberta"
 __snake_case				=						RobertaSeriesConfig



 def __init__(			self	,   _UpperCAmelCase      ):
       super().__init__(_UpperCAmelCase      )
       snake_case_ =   XLMRobertaModel(_UpperCAmelCase      )
       snake_case_ =   nn.Linear(config.hidden_size	,   config.project_dim      )
       snake_case_ =   getattr(_UpperCAmelCase	,   '''has_pre_transformation'''	,   _UpperCAmelCase      )
       if self.has_pre_transformation:
             snake_case_ =   nn.Linear(config.hidden_size	,   config.project_dim      )
             snake_case_ =   nn.LayerNorm(config.hidden_size	,   eps=config.layer_norm_eps      )
       self.post_init()





 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   ):
       snake_case_ =   return_dict if return_dict is not None else self.config.use_return_dict

       snake_case_ =   self.base_model(
           input_ids=_UpperCAmelCase	,   attention_mask=_UpperCAmelCase	,   token_type_ids=_UpperCAmelCase	,   position_ids=_UpperCAmelCase	,   head_mask=_UpperCAmelCase	,   inputs_embeds=_UpperCAmelCase	,   encoder_hidden_states=_UpperCAmelCase	,   encoder_attention_mask=_UpperCAmelCase	,   output_attentions=_UpperCAmelCase	,   output_hidden_states=True if self.has_pre_transformation else output_hidden_states	,   return_dict=_UpperCAmelCase	,   )

       if self.has_pre_transformation:
             snake_case_ =   outputs['''hidden_states'''][-2]
             snake_case_ =   self.pre_LN(_UpperCAmelCase      )
             snake_case_ =   self.transformation_pre(_UpperCAmelCase      )

             return TransformationModelOutput(
                 projection_state=_UpperCAmelCase	,   last_hidden_state=outputs.last_hidden_state	,   hidden_states=outputs.hidden_states	,   attentions=outputs.attentions	,   )
       else:
             snake_case_ =   self.transformation(outputs.last_hidden_state      )
             return TransformationModelOutput(
                 projection_state=_UpperCAmelCase	,   last_hidden_state=outputs.last_hidden_state	,   hidden_states=outputs.hidden_states	,   attentions=outputs.attentions	,   )