





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__           =			logging.get_logger(__name__)

UpperCamelCase__           =			{
    "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 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   Any      =							"""marian"""
							snake_case  :   int      =							["""past_key_values"""]
							snake_case  :   Optional[int]      =							{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






							def __init__(    self	,			__lowerCAmelCase=58101	,			__lowerCAmelCase=None	,			__lowerCAmelCase=1024	,			__lowerCAmelCase=12	,			__lowerCAmelCase=4096	,			__lowerCAmelCase=16	,			__lowerCAmelCase=12	,			__lowerCAmelCase=4096	,			__lowerCAmelCase=16	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase=True	,			__lowerCAmelCase=True	,			__lowerCAmelCase="gelu"	,			__lowerCAmelCase=1024	,			__lowerCAmelCase=0.1	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase=0.02	,			__lowerCAmelCase=58100	,			__lowerCAmelCase=False	,			__lowerCAmelCase=58100	,			__lowerCAmelCase=0	,			__lowerCAmelCase=0	,			__lowerCAmelCase=True	,			**__lowerCAmelCase	,			):
														UpperCamelCase__					  =      vocab_size
														UpperCamelCase__					  =      decoder_vocab_size or vocab_size
														UpperCamelCase__					  =      max_position_embeddings
														UpperCamelCase__					  =      d_model
														UpperCamelCase__					  =      encoder_ffn_dim
														UpperCamelCase__					  =      encoder_layers
														UpperCamelCase__					  =      encoder_attention_heads
														UpperCamelCase__					  =      decoder_ffn_dim
														UpperCamelCase__					  =      decoder_layers
														UpperCamelCase__					  =      decoder_attention_heads
														UpperCamelCase__					  =      dropout
														UpperCamelCase__					  =      attention_dropout
														UpperCamelCase__					  =      activation_dropout
														UpperCamelCase__					  =      activation_function
														UpperCamelCase__					  =      init_std
														UpperCamelCase__					  =      encoder_layerdrop
														UpperCamelCase__					  =      decoder_layerdrop
														UpperCamelCase__					  =      use_cache
														UpperCamelCase__					  =      encoder_layers
														UpperCamelCase__					  =      scale_embedding  # scale factor will be sqrt(d_model) if True
														UpperCamelCase__					  =      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 							__SCREAMING_SNAKE_CASE       (  _a							):






							@property
							# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
							def 					_lowerCamelCase		(    self	):
														if self.task in ["default", "seq2seq-lm"]:
																					UpperCamelCase__					  =      OrderedDict(
																					    [
																					        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
																					        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
																					    ]	)

																					if self.use_past:
																												UpperCamelCase__					  =      {0: """batch"""}
																												UpperCamelCase__					  =      {0: """batch""", 1: """past_decoder_sequence + sequence"""}
																					else:
																												UpperCamelCase__					  =      {0: """batch""", 1: """decoder_sequence"""}
																												UpperCamelCase__					  =      {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__					  =      OrderedDict(
																					    [
																					        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
																					        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
																					    ]	)
																					if self.use_past:
																												UpperCamelCase__  ,  UpperCamelCase__					  =      self.num_layers
																												for i in range(__lowerCAmelCase	):
																																			UpperCamelCase__					  =      {0: """batch""", 2: """past_sequence + sequence"""}
																																			UpperCamelCase__					  =      {0: """batch""", 2: """past_sequence + sequence"""}
														else:
																					UpperCamelCase__					  =      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 					_lowerCamelCase		(    self	):
														if self.task in ["default", "seq2seq-lm"]:
																					UpperCamelCase__					  =      super().outputs
														else:
																					UpperCamelCase__					  =      super(__lowerCAmelCase	,			self	).outputs
																					if self.use_past:
																												UpperCamelCase__  ,  UpperCamelCase__					  =      self.num_layers
																												for i in range(__lowerCAmelCase	):
																																			UpperCamelCase__					  =      {0: """batch""", 2: """past_sequence + sequence"""}
																																			UpperCamelCase__					  =      {0: """batch""", 2: """past_sequence + sequence"""}
														return common_outputs






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = False	,			__lowerCAmelCase = None	,			):
														UpperCamelCase__					  =      self._generate_dummy_inputs_for_encoder_and_decoder(
														    __lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	)

														# Generate decoder inputs
														UpperCamelCase__					  =      seq_length if not self.use_past else 1
														UpperCamelCase__					  =      self._generate_dummy_inputs_for_encoder_and_decoder(
														    __lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	)
														UpperCamelCase__					  =      {f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
														UpperCamelCase__					  =      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__					  =      common_inputs["""input_ids"""].shape
																					UpperCamelCase__					  =      common_inputs["""decoder_input_ids"""].shape[1]
																					UpperCamelCase__  ,  UpperCamelCase__					  =      self.num_attention_heads
																					UpperCamelCase__					  =      (
																					    batch,
																					    num_encoder_attention_heads,
																					    encoder_seq_length,
																					    self._config.hidden_size // num_encoder_attention_heads,
																					)
																					UpperCamelCase__					  =      decoder_seq_length + 3
																					UpperCamelCase__					  =      (
																					    batch,
																					    num_decoder_attention_heads,
																					    decoder_past_length,
																					    self._config.hidden_size // num_decoder_attention_heads,
																					)

																					UpperCamelCase__					  =      torch.cat(
																					    [common_inputs["""decoder_attention_mask"""], torch.ones(__lowerCAmelCase	,			__lowerCAmelCase	)]	,			dim=1	)

																					UpperCamelCase__					  =      []
																					# If the number of encoder and decoder layers are present in the model configuration, both are considered
																					UpperCamelCase__  ,  UpperCamelCase__					  =      self.num_layers
																					UpperCamelCase__					  =      min(__lowerCAmelCase	,			__lowerCAmelCase	)
																					UpperCamelCase__					  =      max(__lowerCAmelCase	,			__lowerCAmelCase	) - min_num_layers
																					UpperCamelCase__					  =      """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__					  =      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 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = False	,			__lowerCAmelCase = None	,			):
														UpperCamelCase__					  =      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__					  =      common_inputs["""input_ids"""].shape
																					# Not using the same length for past_key_values
																					UpperCamelCase__					  =      seqlen + 2
																					UpperCamelCase__  ,  UpperCamelCase__					  =      self.num_layers
																					UpperCamelCase__  ,  UpperCamelCase__					  =      self.num_attention_heads
																					UpperCamelCase__					  =      (
																					    batch,
																					    num_encoder_attention_heads,
																					    past_key_values_length,
																					    self._config.hidden_size // num_encoder_attention_heads,
																					)

																					UpperCamelCase__					  =      common_inputs["""attention_mask"""].dtype
																					UpperCamelCase__					  =      torch.cat(
																					    [common_inputs["""attention_mask"""], torch.ones(__lowerCAmelCase	,			__lowerCAmelCase	,			dtype=__lowerCAmelCase	)]	,			dim=1	)
																					UpperCamelCase__					  =      [
																					    (torch.zeros(__lowerCAmelCase	), torch.zeros(__lowerCAmelCase	)) for _ in range(__lowerCAmelCase	)
																					]
														return common_inputs






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = False	,			__lowerCAmelCase = None	,			):
														# Copied from OnnxConfig.generate_dummy_inputs
														# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
														# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
														UpperCamelCase__					  =      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__					  =      tokenizer.num_special_tokens_to_add(__lowerCAmelCase	)
														UpperCamelCase__					  =      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__					  =      [""" """.join([tokenizer.unk_token]	) * seq_length] * batch_size
														UpperCamelCase__					  =      dict(tokenizer(__lowerCAmelCase	,			return_tensors=__lowerCAmelCase	)	)
														return common_inputs






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = -1	,			__lowerCAmelCase = False	,			__lowerCAmelCase = None	,			):
														if self.task in ["default", "seq2seq-lm"]:
																					UpperCamelCase__					  =      self._generate_dummy_inputs_for_default_and_seqaseq_lm(
																					    __lowerCAmelCase	,			batch_size=__lowerCAmelCase	,			seq_length=__lowerCAmelCase	,			is_pair=__lowerCAmelCase	,			framework=__lowerCAmelCase	)

														else:
																					UpperCamelCase__					  =      self._generate_dummy_inputs_for_causal_lm(
																					    __lowerCAmelCase	,			batch_size=__lowerCAmelCase	,			seq_length=__lowerCAmelCase	,			is_pair=__lowerCAmelCase	,			framework=__lowerCAmelCase	)

														return common_inputs






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







							@property
							def 					_lowerCamelCase		(    self	):
														return 1E-4











import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__           =			logging.get_logger(__name__)

class 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   Optional[Any]      =							"""encoder-decoder"""
							snake_case  :   Optional[int]      =							True






							def __init__(    self	,			**__lowerCAmelCase	):
														super().__init__(**__lowerCAmelCase	)
														assert (
														    "encoder" in kwargs and "decoder" in kwargs
														), "Config has to be initialized with encoder and decoder config"
														UpperCamelCase__					  =      kwargs.pop("""encoder"""	)
														UpperCamelCase__					  =      encoder_config.pop("""model_type"""	)
														UpperCamelCase__					  =      kwargs.pop("""decoder"""	)
														UpperCamelCase__					  =      decoder_config.pop("""model_type"""	)

														from ..auto.configuration_auto import AutoConfig

														UpperCamelCase__					  =      AutoConfig.for_model(__lowerCAmelCase	,			**__lowerCAmelCase	)
														UpperCamelCase__					  =      AutoConfig.for_model(__lowerCAmelCase	,			**__lowerCAmelCase	)
														UpperCamelCase__					  =      True






							@classmethod
							def 					_lowerCamelCase		(    cls	,			__lowerCAmelCase	,			__lowerCAmelCase	,			**__lowerCAmelCase	):
														logger.info("""Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config"""	)
														UpperCamelCase__					  =      True
														UpperCamelCase__					  =      True

														return cls(encoder=encoder_config.to_dict()	,			decoder=decoder_config.to_dict()	,			**__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      copy.deepcopy(self.__dict__	)
														UpperCamelCase__					  =      self.encoder.to_dict()
														UpperCamelCase__					  =      self.decoder.to_dict()
														UpperCamelCase__					  =      self.__class__.model_type
														return output




