





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


_lowerCamelCase 					=      logging.get_logger(__name__)

_lowerCamelCase 					=      {
    '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 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       "marian"
							lowerCamelCase_							     =       ["past_key_values"]
							lowerCamelCase_							     =       {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}


							def __init__(				self     :Tuple				,					__A     :Union[str, Any]=5_8101				,					__A     :Dict=None				,					__A     :str=1024				,					__A     :Optional[int]=12				,					__A     :Optional[int]=4096				,					__A     :Union[str, Any]=16				,					__A     :List[Any]=12				,					__A     :int=4096				,					__A     :Dict=16				,					__A     :Optional[Any]=0.0				,					__A     :Any=0.0				,					__A     :str=True				,					__A     :int=True				,					__A     :Union[str, Any]="gelu"				,					__A     :str=1024				,					__A     :List[Any]=0.1				,					__A     :Tuple=0.0				,					__A     :str=0.0				,					__A     :Dict=0.0_2				,					__A     :Tuple=5_8100				,					__A     :Optional[Any]=False				,					__A     :Dict=5_8100				,					__A     :int=0				,					__A     :str=0				,					__A     :Union[str, Any]=True				,					**__A     :Dict				,					) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   vocab_size
									SCREAMING_SNAKE_CASE__								=   decoder_vocab_size or vocab_size
									SCREAMING_SNAKE_CASE__								=   max_position_embeddings
									SCREAMING_SNAKE_CASE__								=   d_model
									SCREAMING_SNAKE_CASE__								=   encoder_ffn_dim
									SCREAMING_SNAKE_CASE__								=   encoder_layers
									SCREAMING_SNAKE_CASE__								=   encoder_attention_heads
									SCREAMING_SNAKE_CASE__								=   decoder_ffn_dim
									SCREAMING_SNAKE_CASE__								=   decoder_layers
									SCREAMING_SNAKE_CASE__								=   decoder_attention_heads
									SCREAMING_SNAKE_CASE__								=   dropout
									SCREAMING_SNAKE_CASE__								=   attention_dropout
									SCREAMING_SNAKE_CASE__								=   activation_dropout
									SCREAMING_SNAKE_CASE__								=   activation_function
									SCREAMING_SNAKE_CASE__								=   init_std
									SCREAMING_SNAKE_CASE__								=   encoder_layerdrop
									SCREAMING_SNAKE_CASE__								=   decoder_layerdrop
									SCREAMING_SNAKE_CASE__								=   use_cache
									SCREAMING_SNAKE_CASE__								=   encoder_layers
									SCREAMING_SNAKE_CASE__								=   scale_embedding  # scale factor will be sqrt(d_model) if True
									SCREAMING_SNAKE_CASE__								=   share_encoder_decoder_embeddings
									super().__init__(
									    pad_token_id=__A				,					eos_token_id=__A				,					is_encoder_decoder=__A				,					decoder_start_token_id=__A				,					forced_eos_token_id=__A				,					**__A				,					)




class 	UpperCamelCase_							(  UpperCamelCase__     ):


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



									"""simple docstring"""



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

											if self.use_past:
													SCREAMING_SNAKE_CASE__								=   {0: """batch"""}
													SCREAMING_SNAKE_CASE__								=   {0: """batch""", 1: """past_decoder_sequence + sequence"""}
											else:
													SCREAMING_SNAKE_CASE__								=   {0: """batch""", 1: """decoder_sequence"""}
													SCREAMING_SNAKE_CASE__								=   {0: """batch""", 1: """decoder_sequence"""}

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



									"""simple docstring"""



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


							def   _snake_case	(				self     :List[Any]				,					__A     :PreTrainedTokenizer				,					__A     :int = -1				,					__A     :int = -1				,					__A     :bool = False				,					__A     :Optional[TensorType] = None				,					) -> Mapping[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self._generate_dummy_inputs_for_encoder_and_decoder(
									    __A				,					__A				,					__A				,					__A				,					__A		)

									# Generate decoder inputs
									SCREAMING_SNAKE_CASE__								=   seq_length if not self.use_past else 1
									SCREAMING_SNAKE_CASE__								=   self._generate_dummy_inputs_for_encoder_and_decoder(
									    __A				,					__A				,					__A				,					__A				,					__A		)
									SCREAMING_SNAKE_CASE__								=   {f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
									SCREAMING_SNAKE_CASE__								=   dict(**__A				,					**__A		)

									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__								=   common_inputs["""input_ids"""].shape
											SCREAMING_SNAKE_CASE__								=   common_inputs["""decoder_input_ids"""].shape[1]
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self.num_attention_heads
											SCREAMING_SNAKE_CASE__								=   (
											    batch,
											    num_encoder_attention_heads,
											    encoder_seq_length,
											    self._config.hidden_size // num_encoder_attention_heads,
											)
											SCREAMING_SNAKE_CASE__								=   decoder_seq_length + 3
											SCREAMING_SNAKE_CASE__								=   (
											    batch,
											    num_decoder_attention_heads,
											    decoder_past_length,
											    self._config.hidden_size // num_decoder_attention_heads,
											)

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

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

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


							def   _snake_case	(				self     :Dict				,					__A     :PreTrainedTokenizer				,					__A     :int = -1				,					__A     :int = -1				,					__A     :bool = False				,					__A     :Optional[TensorType] = None				,					) -> Mapping[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self._generate_dummy_inputs_for_encoder_and_decoder(
									    __A				,					__A				,					__A				,					__A				,					__A		)

									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__								=   common_inputs["""input_ids"""].shape
											# Not using the same length for past_key_values
											SCREAMING_SNAKE_CASE__								=   seqlen + 2
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self.num_layers
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self.num_attention_heads
											SCREAMING_SNAKE_CASE__								=   (
											    batch,
											    num_encoder_attention_heads,
											    past_key_values_length,
											    self._config.hidden_size // num_encoder_attention_heads,
											)

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


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   compute_effective_axis_dimension(
									    __A				,					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__								=   tokenizer.num_special_tokens_to_add(__A		)
									SCREAMING_SNAKE_CASE__								=   compute_effective_axis_dimension(
									    __A				,					fixed_dimension=OnnxConfig.default_fixed_sequence				,					num_token_to_add=__A		)

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


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



									"""simple docstring"""



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

									else:
											SCREAMING_SNAKE_CASE__								=   self._generate_dummy_inputs_for_causal_lm(
											    __A				,					batch_size=__A				,					seq_length=__A				,					is_pair=__A				,					framework=__A		)

									return common_inputs


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :List[Any]				,					__A     :Tuple				,					__A     :Dict				,					__A     :Any		) -> str:



									"""simple docstring"""



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


							@property
							def   _snake_case	(				self     :Tuple		) -> float:



									"""simple docstring"""



									return 1E-4






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


_lowerCamelCase 					=      {'configuration_speech_encoder_decoder': ['SpeechEncoderDecoderConfig']}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCamelCase 					=      ['SpeechEncoderDecoderModel']

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCamelCase 					=      ['FlaxSpeechEncoderDecoderModel']

if TYPE_CHECKING:
	from .configuration_speech_encoder_decoder import SpeechEncoderDecoderConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_speech_encoder_decoder import SpeechEncoderDecoderModel

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_flax_speech_encoder_decoder import FlaxSpeechEncoderDecoderModel

else:
	import sys

	_lowerCamelCase 					=      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)