




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







						'''simple docstring'''





						def __init__(		self				:	Union[str, Any]			,      *__lowercase				:	str			,      __lowercase				:	List[Any]=None			,      __lowercase				:	Any=None			,      **__lowercase				:	Tuple		):






							"""simple docstring"""






							super().__init__(*__lowercase			,      **__lowercase		)
							snake_case_						       =	eval_examples
							snake_case_						       =	post_process_function






						def 			snake_case__		(		self				:	Any			,      __lowercase				:	Optional[Dataset] = None			,      __lowercase				:	List[Any]=None			,      __lowercase				:	Optional[List[str]] = None			,      __lowercase				:	str = "eval"			,      **__lowercase				:	Tuple			,      ):






							"""simple docstring"""






							snake_case_						       =	gen_kwargs.copy()
							snake_case_						       =	(
							    gen_kwargs["max_length"] if gen_kwargs.get("max_length"		) is not None else self.args.generation_max_length
							)
							snake_case_						       =	(
							    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"		) is not None else self.args.generation_num_beams
							)
							snake_case_						       =	gen_kwargs

							snake_case_						       =	self.eval_dataset if eval_dataset is None else eval_dataset
							snake_case_						       =	self.get_eval_dataloader(__lowercase		)
							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_						       =	time.time()
							snake_case_						       =	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
							try:
								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:
								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(
							        __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
								snake_case_						       =	self.post_process_function(__lowercase			,      __lowercase			,      __lowercase		)
								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}_"		):
										snake_case_						       =	metrics.pop(__lowercase		)

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

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






						def 			snake_case__		(		self				:	Any			,      __lowercase				:	Optional[Any]			,      __lowercase				:	Tuple			,      __lowercase				:	Tuple=None			,      __lowercase				:	str = "test"			,      **__lowercase				:	int		):






							"""simple docstring"""






							snake_case_						       =	gen_kwargs.copy()

							snake_case_						       =	self.get_test_dataloader(__lowercase		)

							# Temporarily disable metric computation, we will do it in the loop here.
							snake_case_						       =	self.compute_metrics
							snake_case_						       =	None
							snake_case_						       =	time.time()
							snake_case_						       =	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
							try:
								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:
								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(
							        __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

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






import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin
class 						UpperCAmelCase  (							UpperCAmelCase__			,  UpperCAmelCase__			,  UpperCAmelCase__     ):







						'''simple docstring'''





						@register_to_config
						def __init__(		self				:	Tuple			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	float			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	int			,      __lowercase				:	str			,      __lowercase				:	bool = False			,      ):






							"""simple docstring"""






							super().__init__()

							snake_case_						       =	nn.Embedding(__lowercase			,      __lowercase		)

							snake_case_						       =	nn.Embedding(__lowercase			,      __lowercase		)
							snake_case_						       =	False

							snake_case_						       =	nn.Dropout(p=__lowercase		)

							snake_case_						       =	TaConfig(
							    vocab_size=__lowercase			,      d_model=__lowercase			,      num_heads=__lowercase			,      d_kv=__lowercase			,      d_ff=__lowercase			,      dropout_rate=__lowercase			,      feed_forward_proj=__lowercase			,      is_decoder=__lowercase			,      is_encoder_decoder=__lowercase			,      )

							snake_case_						       =	nn.ModuleList()
							for lyr_num in range(__lowercase		):
								snake_case_						       =	TaBlock(__lowercase		)
								self.encoders.append(__lowercase		)

							snake_case_						       =	TaLayerNorm(__lowercase		)
							snake_case_						       =	nn.Dropout(p=__lowercase		)






						def 			snake_case__		(		self				:	str			,      __lowercase				:	Dict			,      __lowercase				:	str		):






							"""simple docstring"""






							snake_case_						       =	self.token_embedder(__lowercase		)

							snake_case_						       =	encoder_input_tokens.shape[1]
							snake_case_						       =	torch.arange(__lowercase			,      device=encoder_input_tokens.device		)
							x += self.position_encoding(__lowercase		)

							snake_case_						       =	self.dropout_pre(__lowercase		)

							# inverted the attention mask
							snake_case_						       =	encoder_input_tokens.size()
							snake_case_						       =	self.get_extended_attention_mask(__lowercase			,      __lowercase		)

							for lyr in self.encoders:
								snake_case_						       =	lyr(__lowercase			,      __lowercase		)[0]
							snake_case_						       =	self.layer_norm(__lowercase		)

							return self.dropout_post(__lowercase		), encoder_inputs_mask
