





"""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


a						: int     	=			logging.get_logger(__name__)

a						: str     	=			{
    '''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				(						a__						):
							lowerCamelCase   :					Any      					="""marian"""
							lowerCamelCase   :					Any      					=["""past_key_values"""]
							lowerCamelCase   :					Optional[int]      					={"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






							def __init__(   self      ,					lowerCAmelCase__=5_8101      ,					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__=5_8100      ,					lowerCAmelCase__=False      ,					lowerCAmelCase__=5_8100      ,					lowerCAmelCase__=0      ,					lowerCAmelCase__=0      ,					lowerCAmelCase__=True      ,					**lowerCAmelCase__      ,					)	-> Dict:
								a      :		Optional[Any]	=				vocab_size
								a      :		Dict	=				decoder_vocab_size or vocab_size
								a      :		Dict	=				max_position_embeddings
								a      :		Dict	=				d_model
								a      :		Optional[Any]	=				encoder_ffn_dim
								a      :		Union[str, Any]	=				encoder_layers
								a      :		Any	=				encoder_attention_heads
								a      :		Optional[Any]	=				decoder_ffn_dim
								a      :		Optional[int]	=				decoder_layers
								a      :		Optional[int]	=				decoder_attention_heads
								a      :		Dict	=				dropout
								a      :		Optional[Any]	=				attention_dropout
								a      :		Optional[int]	=				activation_dropout
								a      :		int	=				activation_function
								a      :		Union[str, Any]	=				init_std
								a      :		Optional[int]	=				encoder_layerdrop
								a      :		Dict	=				decoder_layerdrop
								a      :		Any	=				use_cache
								a      :		List[str]	=				encoder_layers
								a      :		str	=				scale_embedding  # scale factor will be sqrt(d_model) if True
								a      :		List[str]	=				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        __UpperCamelCase				(						a__						):






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

									if self.use_past:
										a      :		Tuple	=				{0: "batch"}
										a      :		Tuple	=				{0: "batch", 1: "past_decoder_sequence + sequence"}
									else:
										a      :		str	=				{0: "batch", 1: "decoder_sequence"}
										a      :		List[Any]	=				{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.
									a      :		int	=				OrderedDict(
									    [
									        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
									        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
									    ]      )
									if self.use_past:
										a,       a      :		Tuple	=				self.num_layers
										for i in range(lowerCAmelCase__      ):
											a      :		Optional[Any]	=				{0: "batch", 2: "past_sequence + sequence"}
											a      :		Optional[Any]	=				{0: "batch", 2: "past_sequence + sequence"}
								else:
									a      :		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     __a   (   self      )	-> Mapping[str, Mapping[int, str]]:
								if self.task in ["default", "seq2seq-lm"]:
									a      :		int	=				super().outputs
								else:
									a      :		int	=				super(lowerCAmelCase__      ,					self      ).outputs
									if self.use_past:
										a,       a      :		str	=				self.num_layers
										for i in range(lowerCAmelCase__      ):
											a      :		Union[str, Any]	=				{0: "batch", 2: "past_sequence + sequence"}
											a      :		List[Any]	=				{0: "batch", 2: "past_sequence + sequence"}
								return common_outputs






							def     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = False      ,					lowerCAmelCase__ = None      ,					)	-> Mapping[str, Any]:
								a      :		List[Any]	=				self._generate_dummy_inputs_for_encoder_and_decoder(
								    lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      )

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

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

									a      :		Optional[Any]	=				[]
									# If the number of encoder and decoder layers are present in the model configuration, both are considered
									a,       a      :		Union[str, Any]	=				self.num_layers
									a      :		Tuple	=				min(lowerCAmelCase__      ,					lowerCAmelCase__      )
									a      :		Tuple	=				max(lowerCAmelCase__      ,					lowerCAmelCase__      ) - min_num_layers
									a      :		Tuple	=				"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.
									a      :		Any	=				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     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = False      ,					lowerCAmelCase__ = None      ,					)	-> Mapping[str, Any]:
								a      :		List[str]	=				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
									a,       a      :		int	=				common_inputs["input_ids"].shape
									# Not using the same length for past_key_values
									a      :		Optional[Any]	=				seqlen + 2
									a,       a      :		List[str]	=				self.num_layers
									a,       a      :		Optional[int]	=				self.num_attention_heads
									a      :		Dict	=				(
									    batch,
									    num_encoder_attention_heads,
									    past_key_values_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)

									a      :		Optional[Any]	=				common_inputs["attention_mask"].dtype
									a      :		Any	=				torch.cat(
									    [common_inputs["attention_mask"], torch.ones(lowerCAmelCase__      ,					lowerCAmelCase__      ,					dtype=lowerCAmelCase__      )]      ,					dim=1      )
									a      :		Dict	=				[
									    (torch.zeros(lowerCAmelCase__      ), torch.zeros(lowerCAmelCase__      )) for _ in range(lowerCAmelCase__      )
									]
								return common_inputs






							def     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = False      ,					lowerCAmelCase__ = None      ,					)	-> Mapping[str, Any]:
								# 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
								a      :		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
								a      :		Optional[Any]	=				tokenizer.num_special_tokens_to_add(lowerCAmelCase__      )
								a      :		str	=				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
								a      :		Optional[int]	=				[" ".join([tokenizer.unk_token]      ) * seq_length] * batch_size
								a      :		Optional[Any]	=				dict(tokenizer(lowerCAmelCase__      ,					return_tensors=lowerCAmelCase__      )      )
								return common_inputs






							def     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = -1      ,					lowerCAmelCase__ = False      ,					lowerCAmelCase__ = None      ,					)	-> Mapping[str, Any]:
								if self.task in ["default", "seq2seq-lm"]:
									a      :		Tuple	=				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
									    lowerCAmelCase__      ,					batch_size=lowerCAmelCase__      ,					seq_length=lowerCAmelCase__      ,					is_pair=lowerCAmelCase__      ,					framework=lowerCAmelCase__      )

								else:
									a      :		List[str]	=				self._generate_dummy_inputs_for_causal_lm(
									    lowerCAmelCase__      ,					batch_size=lowerCAmelCase__      ,					seq_length=lowerCAmelCase__      ,					is_pair=lowerCAmelCase__      ,					framework=lowerCAmelCase__      )

								return common_inputs






							def     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      )	-> Dict:
								if self.task in ["default", "seq2seq-lm"]:
									a      :		List[str]	=				super()._flatten_past_key_values_(lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      )
								else:
									a      :		Optional[Any]	=				super(lowerCAmelCase__      ,					self      )._flatten_past_key_values_(
									    lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      )






							@property
							def     __a   (   self      )	-> float:
								return 1E-4









"""simple docstring"""




import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin


class        __UpperCamelCase				(						a__		,       a__						):






							@register_to_config
							def __init__(   self      ,					*,
							    lowerCAmelCase__ = 4      ,					lowerCAmelCase__ = 768      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					)	-> Optional[Any]:
								super().__init__()

								a      :		Tuple	=				nn.Parameter(torch.zeros(lowerCAmelCase__      )      )

								# parameters for additional clip time embeddings
								a      :		str	=				nn.Linear(lowerCAmelCase__      ,					lowerCAmelCase__      )
								a      :		Any	=				nn.Linear(lowerCAmelCase__      ,					lowerCAmelCase__      )

								# parameters for encoder hidden states
								a      :		int	=				clip_extra_context_tokens
								a      :		int	=				nn.Linear(
								    lowerCAmelCase__      ,					self.clip_extra_context_tokens * cross_attention_dim      )
								a      :		Any	=				nn.Linear(lowerCAmelCase__      ,					lowerCAmelCase__      )
								a      :		str	=				nn.LayerNorm(lowerCAmelCase__      )






							def     __a   (   self      ,					*, lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      )	-> Union[str, Any]:
								if do_classifier_free_guidance:
									# Add the classifier free guidance embeddings to the image embeddings
									a      :		str	=				image_embeddings.shape[0]
									a      :		Optional[int]	=				self.learned_classifier_free_guidance_embeddings.unsqueeze(0      )
									a      :		Any	=				classifier_free_guidance_embeddings.expand(
									    lowerCAmelCase__      ,					-1      )
									a      :		Any	=				torch.cat([classifier_free_guidance_embeddings, image_embeddings]      ,					dim=0      )

								# The image embeddings batch size and the text embeddings batch size are equal
								assert image_embeddings.shape[0] == prompt_embeds.shape[0]

								a      :		List[str]	=				prompt_embeds.shape[0]

								# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
								# adding CLIP embeddings to the existing timestep embedding, ...
								a      :		Dict	=				self.embedding_proj(lowerCAmelCase__      )
								a      :		List[str]	=				self.clip_image_embeddings_project_to_time_embeddings(lowerCAmelCase__      )
								a      :		Dict	=				time_projected_image_embeddings + time_projected_prompt_embeds

								# ... and by projecting CLIP embeddings into four
								# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
								a      :		Union[str, Any]	=				self.clip_extra_context_tokens_proj(lowerCAmelCase__      )
								a      :		List[str]	=				clip_extra_context_tokens.reshape(lowerCAmelCase__      ,					-1      ,					self.clip_extra_context_tokens      )
								a      :		Optional[Any]	=				clip_extra_context_tokens.permute(0      ,					2      ,					1      )

								a      :		Optional[int]	=				self.encoder_hidden_states_proj(lowerCAmelCase__      )
								a      :		str	=				self.text_encoder_hidden_states_norm(lowerCAmelCase__      )
								a      :		List[str]	=				torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]      ,					dim=1      )

								return text_encoder_hidden_states, additive_clip_time_embeddings


