


from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


__UpperCamelCase				:				List[str]  				= logging.get_logger(__name__)

__UpperCamelCase				:				Optional[int]  				= {
    'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json',
    # See all Marian models at https://huggingface.co/models?filter=marian
}

class 					__magic_name__		(   A__):
     A:						Tuple								=					"marian"
     A:						List[Any]								=					["past_key_values"]
     A:						Optional[int]								=					{"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}


     def __init__(					self							:   List[str]							,		lowerCamelCase__							:   List[str]=58101							,		lowerCamelCase__							:   List[str]=None							,		lowerCamelCase__							:   int=1024							,		lowerCamelCase__							:   Union[str, Any]=12							,		lowerCamelCase__							:   int=4096							,		lowerCamelCase__							:   List[str]=16							,		lowerCamelCase__							:   str=12							,		lowerCamelCase__							:   List[str]=4096							,		lowerCamelCase__							:   List[Any]=16							,		lowerCamelCase__							:   Dict=0.0							,		lowerCamelCase__							:   Optional[Any]=0.0							,		lowerCamelCase__							:   int=True							,		lowerCamelCase__							:   int=True							,		lowerCamelCase__							:   str="gelu"							,		lowerCamelCase__							:   str=1024							,		lowerCamelCase__							:   Any=0.1							,		lowerCamelCase__							:   Tuple=0.0							,		lowerCamelCase__							:   List[str]=0.0							,		lowerCamelCase__							:   List[str]=0.02							,		lowerCamelCase__							:   List[str]=58100							,		lowerCamelCase__							:   Dict=False							,		lowerCamelCase__							:   List[Any]=58100							,		lowerCamelCase__							:   Optional[Any]=0							,		lowerCamelCase__							:   Any=0							,		lowerCamelCase__							:   str=True							,		**lowerCamelCase__							:   Union[str, Any]							,		)     ->					List[str]:





      '''simple docstring'''






      UpperCamelCase__    :			Union[str, Any]      				=      vocab_size
      UpperCamelCase__    :			Tuple      				=      decoder_vocab_size or vocab_size
      UpperCamelCase__    :			str      				=      max_position_embeddings
      UpperCamelCase__    :			Optional[Any]      				=      d_model
      UpperCamelCase__    :			Any      				=      encoder_ffn_dim
      UpperCamelCase__    :			Any      				=      encoder_layers
      UpperCamelCase__    :			Union[str, Any]      				=      encoder_attention_heads
      UpperCamelCase__    :			int      				=      decoder_ffn_dim
      UpperCamelCase__    :			Tuple      				=      decoder_layers
      UpperCamelCase__    :			List[str]      				=      decoder_attention_heads
      UpperCamelCase__    :			Any      				=      dropout
      UpperCamelCase__    :			Optional[Any]      				=      attention_dropout
      UpperCamelCase__    :			Optional[int]      				=      activation_dropout
      UpperCamelCase__    :			List[str]      				=      activation_function
      UpperCamelCase__    :			Any      				=      init_std
      UpperCamelCase__    :			Union[str, Any]      				=      encoder_layerdrop
      UpperCamelCase__    :			Optional[Any]      				=      decoder_layerdrop
      UpperCamelCase__    :			Optional[Any]      				=      use_cache
      UpperCamelCase__    :			List[str]      				=      encoder_layers
      UpperCamelCase__    :			Dict      				=      scale_embedding  # scale factor will be sqrt(d_model) if True
      UpperCamelCase__    :			Optional[Any]      				=      share_encoder_decoder_embeddings
      super().__init__(
          pad_token_id=lowerCamelCase__							,		eos_token_id=lowerCamelCase__							,		is_encoder_decoder=lowerCamelCase__							,		decoder_start_token_id=lowerCamelCase__							,		forced_eos_token_id=lowerCamelCase__							,		**lowerCamelCase__							,		)



class 					__magic_name__		(   A__):


     @property
     # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
     def    UpperCAmelCase__		(					self							:   Any	)     ->					Optional[int]:





      '''simple docstring'''






      if self.task in ["default", "seq2seq-lm"]:
       UpperCamelCase__    :			List[Any]      				=      OrderedDict(
           [
               ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
           ]	)

       if self.use_past:
        UpperCamelCase__    :			Tuple      				=      {0: '''batch'''}
        UpperCamelCase__    :			Dict      				=      {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
       else:
        UpperCamelCase__    :			Tuple      				=      {0: '''batch''', 1: '''decoder_sequence'''}
        UpperCamelCase__    :			Dict      				=      {0: '''batch''', 1: '''decoder_sequence'''}

       if self.use_past:
        self.fill_with_past_key_values_(lowerCamelCase__							,		direction='''inputs'''	)
      elif self.task == "causal-lm":
       # TODO: figure this case out.
       UpperCamelCase__    :			Dict      				=      OrderedDict(
           [
               ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
           ]	)
       if self.use_past:
        UpperCamelCase__				, UpperCamelCase__    :			Dict      				=      self.num_layers
        for i in range(lowerCamelCase__	):
         UpperCamelCase__    :			Optional[Any]      				=      {0: '''batch''', 2: '''past_sequence + sequence'''}
         UpperCamelCase__    :			Tuple      				=      {0: '''batch''', 2: '''past_sequence + sequence'''}
      else:
       UpperCamelCase__    :			List[Any]      				=      OrderedDict(
           [
               ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
               ('''decoder_input_ids''', {0: '''batch''', 1: '''decoder_sequence'''}),
               ('''decoder_attention_mask''', {0: '''batch''', 1: '''decoder_sequence'''}),
           ]	)

      return common_inputs


     @property
     # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
     def    UpperCAmelCase__		(					self							:   List[str]	)     ->					str:





      '''simple docstring'''






      if self.task in ["default", "seq2seq-lm"]:
       UpperCamelCase__    :			int      				=      super().outputs
      else:
       UpperCamelCase__    :			Tuple      				=      super(lowerCamelCase__							,		self	).outputs
       if self.use_past:
        UpperCamelCase__				, UpperCamelCase__    :			List[Any]      				=      self.num_layers
        for i in range(lowerCamelCase__	):
         UpperCamelCase__    :			List[str]      				=      {0: '''batch''', 2: '''past_sequence + sequence'''}
         UpperCamelCase__    :			Dict      				=      {0: '''batch''', 2: '''past_sequence + sequence'''}
      return common_outputs


     def    UpperCAmelCase__		(					self							:   str							,		lowerCamelCase__							:   PreTrainedTokenizer							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   bool = False							,		lowerCamelCase__							:   Optional[TensorType] = None							,		)     ->					int:





      '''simple docstring'''






      UpperCamelCase__    :			Optional[Any]      				=      self._generate_dummy_inputs_for_encoder_and_decoder(
          lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__	)

      # Generate decoder inputs
      UpperCamelCase__    :			Tuple      				=      seq_length if not self.use_past else 1
      UpperCamelCase__    :			List[str]      				=      self._generate_dummy_inputs_for_encoder_and_decoder(
          lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__	)
      UpperCamelCase__    :			Dict      				=      {F"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
      UpperCamelCase__    :			str      				=      dict(**lowerCamelCase__							,		**lowerCamelCase__	)

      if self.use_past:
       if not is_torch_available():
        raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''	)
       else:
        import torch
       UpperCamelCase__				, UpperCamelCase__    :			Any      				=      common_inputs['''input_ids'''].shape
       UpperCamelCase__    :			Union[str, Any]      				=      common_inputs['''decoder_input_ids'''].shape[1]
       UpperCamelCase__				, UpperCamelCase__    :			Union[str, Any]      				=      self.num_attention_heads
       UpperCamelCase__    :			Optional[Any]      				=      (
           batch,
           num_encoder_attention_heads,
           encoder_seq_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )
       UpperCamelCase__    :			Optional[Any]      				=      decoder_seq_length + 3
       UpperCamelCase__    :			Any      				=      (
           batch,
           num_decoder_attention_heads,
           decoder_past_length,
           self._config.hidden_size // num_decoder_attention_heads,
       )

       UpperCamelCase__    :			int      				=      torch.cat(
           [common_inputs['''decoder_attention_mask'''], torch.ones(lowerCamelCase__							,		lowerCamelCase__	)]							,		dim=1	)

       UpperCamelCase__    :			Any      				=      []
       # If the number of encoder and decoder layers are present in the model configuration, both are considered
       UpperCamelCase__				, UpperCamelCase__    :			Optional[int]      				=      self.num_layers
       UpperCamelCase__    :			Optional[Any]      				=      min(lowerCamelCase__							,		lowerCamelCase__	)
       UpperCamelCase__    :			List[Any]      				=      max(lowerCamelCase__							,		lowerCamelCase__	) - min_num_layers
       UpperCamelCase__    :			str      				=      '''encoder''' if num_encoder_layers > num_decoder_layers else '''decoder'''

       for _ in range(lowerCamelCase__	):
        common_inputs["past_key_values"].append(
            (
                torch.zeros(lowerCamelCase__	),
                torch.zeros(lowerCamelCase__	),
                torch.zeros(lowerCamelCase__	),
                torch.zeros(lowerCamelCase__	),
            )	)
       # TODO: test this.
       UpperCamelCase__    :			str      				=      encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
       for _ in range(lowerCamelCase__							,		lowerCamelCase__	):
        common_inputs["past_key_values"].append((torch.zeros(lowerCamelCase__	), torch.zeros(lowerCamelCase__	))	)
      return common_inputs


     def    UpperCAmelCase__		(					self							:   Optional[Any]							,		lowerCamelCase__							:   PreTrainedTokenizer							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   bool = False							,		lowerCamelCase__							:   Optional[TensorType] = None							,		)     ->					Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__    :			List[Any]      				=      self._generate_dummy_inputs_for_encoder_and_decoder(
          lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__	)

      if self.use_past:
       if not is_torch_available():
        raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''	)
       else:
        import torch
       UpperCamelCase__				, UpperCamelCase__    :			List[Any]      				=      common_inputs['''input_ids'''].shape
       # Not using the same length for past_key_values
       UpperCamelCase__    :			Dict      				=      seqlen + 2
       UpperCamelCase__				, UpperCamelCase__    :			Dict      				=      self.num_layers
       UpperCamelCase__				, UpperCamelCase__    :			Dict      				=      self.num_attention_heads
       UpperCamelCase__    :			List[Any]      				=      (
           batch,
           num_encoder_attention_heads,
           past_key_values_length,
           self._config.hidden_size // num_encoder_attention_heads,
       )

       UpperCamelCase__    :			Optional[Any]      				=      common_inputs['''attention_mask'''].dtype
       UpperCamelCase__    :			List[Any]      				=      torch.cat(
           [common_inputs['''attention_mask'''], torch.ones(lowerCamelCase__							,		lowerCamelCase__							,		dtype=lowerCamelCase__	)]							,		dim=1	)
       UpperCamelCase__    :			Any      				=      [
           (torch.zeros(lowerCamelCase__	), torch.zeros(lowerCamelCase__	)) for _ in range(lowerCamelCase__	)
       ]
      return common_inputs


     def    UpperCAmelCase__		(					self							:   List[Any]							,		lowerCamelCase__							:   PreTrainedTokenizer							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   bool = False							,		lowerCamelCase__							:   Optional[TensorType] = None							,		)     ->					int:





      '''simple docstring'''






      UpperCamelCase__    :			Tuple      				=      compute_effective_axis_dimension(
          lowerCamelCase__							,		fixed_dimension=OnnxConfig.default_fixed_batch							,		num_token_to_add=0	)

      # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
      UpperCamelCase__    :			int      				=      tokenizer.num_special_tokens_to_add(lowerCamelCase__	)
      UpperCamelCase__    :			Any      				=      compute_effective_axis_dimension(
          lowerCamelCase__							,		fixed_dimension=OnnxConfig.default_fixed_sequence							,		num_token_to_add=lowerCamelCase__	)

      # Generate dummy inputs according to compute batch and sequence
      UpperCamelCase__    :			Union[str, Any]      				=      [''' '''.join([tokenizer.unk_token]	) * seq_length] * batch_size
      UpperCamelCase__    :			int      				=      dict(tokenizer(lowerCamelCase__							,		return_tensors=lowerCamelCase__	)	)
      return common_inputs


     def    UpperCAmelCase__		(					self							:   Any							,		lowerCamelCase__							:   PreTrainedTokenizer							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   int = -1							,		lowerCamelCase__							:   bool = False							,		lowerCamelCase__							:   Optional[TensorType] = None							,		)     ->					str:





      '''simple docstring'''






      if self.task in ["default", "seq2seq-lm"]:
       UpperCamelCase__    :			int      				=      self._generate_dummy_inputs_for_default_and_seqaseq_lm(
           lowerCamelCase__							,		batch_size=lowerCamelCase__							,		seq_length=lowerCamelCase__							,		is_pair=lowerCamelCase__							,		framework=lowerCamelCase__	)

      else:
       UpperCamelCase__    :			Optional[int]      				=      self._generate_dummy_inputs_for_causal_lm(
           lowerCamelCase__							,		batch_size=lowerCamelCase__							,		seq_length=lowerCamelCase__							,		is_pair=lowerCamelCase__							,		framework=lowerCamelCase__	)

      return common_inputs


     def    UpperCAmelCase__		(					self							:   List[Any]							,		lowerCamelCase__							:   List[Any]							,		lowerCamelCase__							:   Tuple							,		lowerCamelCase__							:   int							,		lowerCamelCase__							:   Tuple	)     ->					List[str]:





      '''simple docstring'''






      if self.task in ["default", "seq2seq-lm"]:
       UpperCamelCase__    :			str      				=      super()._flatten_past_key_values_(lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__	)
      else:
       UpperCamelCase__    :			str      				=      super(lowerCamelCase__							,		self	)._flatten_past_key_values_(
           lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__							,		lowerCamelCase__	)


     @property
     def    UpperCAmelCase__		(					self							:   Tuple	)     ->					Optional[int]:





      '''simple docstring'''






      return 1E-4










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						:    Dict,  *lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    Any=None,  **lowerCamelCase						:    str  ):

								'''simple docstring'''
								super().__init__(*lowerCamelCase,  **lowerCamelCase  )
								lowercase__    =			eval_examples
								lowercase__    =			post_process_function




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    str=None,  lowerCamelCase						:    Optional[Any]=None,  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    str = "eval"  ):

								'''simple docstring'''
								lowercase__    =			self.eval_dataset if eval_dataset is None else eval_dataset
								lowercase__    =			self.get_eval_dataloader(lowerCamelCase  )
								lowercase__    =			self.eval_examples if eval_examples is None else eval_examples

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

								lowercase__    =			self.callback_handler.on_evaluate(self.args,  self.state,  self.control,  lowerCamelCase  )
								return metrics




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    Any,  lowerCamelCase						:    Dict,  lowerCamelCase						:    int=None,  lowerCamelCase						:    str = "test"  ):

								'''simple docstring'''
								lowercase__    =			self.get_test_dataloader(lowerCamelCase  )

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

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




