
"""simple docstring"""




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

UpperCAmelCase__   :					Dict												=				{
    '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 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				__UpperCamelCase   :		List[str]       =     '''marian'''
				__UpperCamelCase   :		str       =     ['''past_key_values''']
				__UpperCamelCase   :		Optional[Any]       =     {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}



				def __init__(self     , SCREAMING_SNAKE_CASE__=5_81_01     , SCREAMING_SNAKE_CASE__=None     , SCREAMING_SNAKE_CASE__=10_24     , SCREAMING_SNAKE_CASE__=12     , SCREAMING_SNAKE_CASE__=40_96     , SCREAMING_SNAKE_CASE__=16     , SCREAMING_SNAKE_CASE__=12     , SCREAMING_SNAKE_CASE__=40_96     , SCREAMING_SNAKE_CASE__=16     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__="gelu"     , SCREAMING_SNAKE_CASE__=10_24     , SCREAMING_SNAKE_CASE__=0.1     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=0.02     , SCREAMING_SNAKE_CASE__=5_81_00     , SCREAMING_SNAKE_CASE__=False     , SCREAMING_SNAKE_CASE__=5_81_00     , SCREAMING_SNAKE_CASE__=0     , SCREAMING_SNAKE_CASE__=0     , SCREAMING_SNAKE_CASE__=True     , **SCREAMING_SNAKE_CASE__     , )      ->	Union[str, Any]:



											"""simple docstring"""




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






class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				@property
				# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
				def 			__magic_name__					(self					)      ->	Mapping[str, Mapping[int, str]]:



											"""simple docstring"""




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

																		if self.use_past:
																									SCREAMING_SNAKE_CASE__		: str     =     {0: """batch"""}
																									SCREAMING_SNAKE_CASE__		: Any     =     {0: """batch""", 1: """past_decoder_sequence + sequence"""}
																		else:
																									SCREAMING_SNAKE_CASE__		: List[Any]     =     {0: """batch""", 1: """decoder_sequence"""}
																									SCREAMING_SNAKE_CASE__		: Tuple     =     {0: """batch""", 1: """decoder_sequence"""}

																		if self.use_past:
																									self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE__     , direction="""inputs"""					)
											elif self.task == "causal-lm":
																		# TODO: figure this case out.
																		SCREAMING_SNAKE_CASE__		: Optional[Any]     =     OrderedDict(
																		    [
																		        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
																		        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
																		    ]					)
																		if self.use_past:
																									SCREAMING_SNAKE_CASE__   ,   SCREAMING_SNAKE_CASE__		: Dict     =     self.num_layers
																									for i in range(SCREAMING_SNAKE_CASE__					):
																																SCREAMING_SNAKE_CASE__		: List[Any]     =     {0: """batch""", 2: """past_sequence + sequence"""}
																																SCREAMING_SNAKE_CASE__		: Optional[Any]     =     {0: """batch""", 2: """past_sequence + sequence"""}
											else:
																		SCREAMING_SNAKE_CASE__		: Optional[int]     =     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 			__magic_name__					(self					)      ->	Mapping[str, Mapping[int, str]]:



											"""simple docstring"""




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



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = False     , SCREAMING_SNAKE_CASE__ = None     , )      ->	Mapping[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Tuple     =     self._generate_dummy_inputs_for_encoder_and_decoder(
											    SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)

											# Generate decoder inputs
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     seq_length if not self.use_past else 1
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     self._generate_dummy_inputs_for_encoder_and_decoder(
											    SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Tuple     =     {F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
											SCREAMING_SNAKE_CASE__		: int     =     dict(**SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)

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

																		SCREAMING_SNAKE_CASE__		: Optional[int]     =     torch.cat(
																		    [common_inputs["""decoder_attention_mask"""], torch.ones(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)]     , dim=1					)

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

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



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = False     , SCREAMING_SNAKE_CASE__ = None     , )      ->	Mapping[str, Any]:



											"""simple docstring"""




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

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

																		SCREAMING_SNAKE_CASE__		: List[str]     =     common_inputs["""attention_mask"""].dtype
																		SCREAMING_SNAKE_CASE__		: Dict     =     torch.cat(
																		    [common_inputs["""attention_mask"""], torch.ones(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , dtype=SCREAMING_SNAKE_CASE__					)]     , dim=1					)
																		SCREAMING_SNAKE_CASE__		: List[str]     =     [
																		    (torch.zeros(SCREAMING_SNAKE_CASE__					), torch.zeros(SCREAMING_SNAKE_CASE__					)) for _ in range(SCREAMING_SNAKE_CASE__					)
																		]
											return common_inputs



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = False     , SCREAMING_SNAKE_CASE__ = None     , )      ->	Mapping[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: str     =     compute_effective_axis_dimension(
											    SCREAMING_SNAKE_CASE__     , 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
											SCREAMING_SNAKE_CASE__		: Tuple     =     tokenizer.num_special_tokens_to_add(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     compute_effective_axis_dimension(
											    SCREAMING_SNAKE_CASE__     , fixed_dimension=OnnxConfig.default_fixed_sequence     , num_token_to_add=SCREAMING_SNAKE_CASE__					)

											# Generate dummy inputs according to compute batch and sequence
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     [""" """.join([tokenizer.unk_token]					) * seq_length] * batch_size
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     dict(tokenizer(SCREAMING_SNAKE_CASE__     , return_tensors=SCREAMING_SNAKE_CASE__					)					)
											return common_inputs



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = -1     , SCREAMING_SNAKE_CASE__ = False     , SCREAMING_SNAKE_CASE__ = None     , )      ->	Mapping[str, Any]:



											"""simple docstring"""




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

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

											return common_inputs



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)      ->	Dict:



											"""simple docstring"""




											if self.task in ["default", "seq2seq-lm"]:
																		SCREAMING_SNAKE_CASE__		: Optional[int]     =     super()._flatten_past_key_values_(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)
											else:
																		SCREAMING_SNAKE_CASE__		: List[str]     =     super(SCREAMING_SNAKE_CASE__     , self					)._flatten_past_key_values_(
																		    SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)



				@property
				def 			__magic_name__					(self					)      ->	float:



											"""simple docstring"""




											return 1E-4






"""simple docstring"""




import torch

def lowercase_						(     ):
							if torch.cuda.is_available():
														SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.cuda.device_count()
							else:
														SCREAMING_SNAKE_CASE__		: int     =     0
							print(f'''Successfully ran on {num_gpus} GPUs'''		)


if __name__ == "__main__":
					main()




