




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


_SCREAMING_SNAKE_CASE							:     Tuple   	=				logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE							:     Optional[Any]   	=				{
    "google/bigbird-pegasus-large-arxiv": (
        "https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json"
    ),
    "google/bigbird-pegasus-large-pubmed": (
        "https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json"
    ),
    "google/bigbird-pegasus-large-bigpatent": (
        "https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json"
    ),
    # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
}




class 				_snake_case	(			lowercase_							):
			lowerCAmelCase_  :   Optional[Any]		 =					"bigbird_pegasus"
			lowerCAmelCase_  :   str		 =					["past_key_values"]
			lowerCAmelCase_  :   List[Any]		 =					{
			    "num_attention_heads": "encoder_attention_heads",
			    "hidden_size": "d_model",
			    "attention_probs_dropout_prob": "attention_dropout",
			}




			def __init__(      self       ,							a__=96_103       ,							a__=4_096       ,							a__=16       ,							a__=4_096       ,							a__=16       ,							a__=16       ,							a__=4_096       ,							a__=16       ,							a__=0.0       ,							a__=0.0       ,							a__=True       ,							a__=True       ,							a__="gelu_new"       ,							a__=1_024       ,							a__=0.1       ,							a__=0.0       ,							a__=0.0       ,							a__=0.0_2       ,							a__=2       ,							a__=0.0       ,							a__=True       ,							a__=0       ,							a__=2       ,							a__=1       ,							a__="block_sparse"       ,							a__=64       ,							a__=3       ,							a__=False       ,							**a__       ,							)       ->       Optional[int]:

										'''simple docstring'''




										snake_case_			      =   vocab_size
										snake_case_			      =   max_position_embeddings
										snake_case_			      =   d_model
										snake_case_			      =   encoder_ffn_dim
										snake_case_			      =   encoder_layers
										snake_case_			      =   encoder_attention_heads
										snake_case_			      =   decoder_ffn_dim
										snake_case_			      =   decoder_layers
										snake_case_			      =   decoder_attention_heads
										snake_case_			      =   dropout
										snake_case_			      =   attention_dropout
										snake_case_			      =   activation_dropout
										snake_case_			      =   activation_function
										snake_case_			      =   init_std
										snake_case_			      =   encoder_layerdrop
										snake_case_			      =   decoder_layerdrop
										snake_case_			      =   classifier_dropout
										snake_case_			      =   use_cache
										snake_case_			      =   encoder_layers
										snake_case_			      =   scale_embedding  # scale factor will be sqrt(d_model) if True

										# extra config
										snake_case_			      =   attention_type
										snake_case_			      =   block_size
										snake_case_			      =   num_random_blocks
										snake_case_			      =   use_bias

										super().__init__(
										    pad_token_id=a__       ,							bos_token_id=a__       ,							eos_token_id=a__       ,							is_encoder_decoder=a__       ,							decoder_start_token_id=a__       ,							**a__       ,							)




class 				_snake_case	(			lowercase_							):




			@property
			def      lowerCAmelCase__       (      self							)       ->       Mapping[str, Mapping[int, str]]:

										'''simple docstring'''




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

																	if self.use_past:
																								snake_case_			      =   {0: "batch"}
																								snake_case_			      =   {0: "batch", 1: "past_decoder_sequence + sequence"}
																	else:
																								snake_case_			      =   {0: "batch", 1: "decoder_sequence"}
																								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.
																	snake_case_			      =   OrderedDict(
																	    [
																	        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
																	        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
																	    ]							)
																	if self.use_past:
																								snake_case_      ,   snake_case_			      =   self.num_layers
																								for i in range(a__							):
																															snake_case_			      =   {0: "batch", 2: "past_sequence + sequence"}
																															snake_case_			      =   {0: "batch", 2: "past_sequence + sequence"}
										else:
																	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
			def      lowerCAmelCase__       (      self							)       ->       Mapping[str, Mapping[int, str]]:

										'''simple docstring'''




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




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

										'''simple docstring'''




										snake_case_			      =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a__       ,							a__       ,							a__       ,							a__       ,							a__							)

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

																	snake_case_			      =   torch.cat(
																	    [common_inputs["decoder_attention_mask"], torch.ones(a__       ,							a__							)]       ,							dim=1							)

																	snake_case_			      =   []
																	# If the number of encoder and decoder layers are present in the model configuration, both are considered
																	snake_case_      ,   snake_case_			      =   self.num_layers
																	snake_case_			      =   min(a__       ,							a__							)
																	snake_case_			      =   max(a__       ,							a__							) - min_num_layers
																	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.
																	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      lowerCAmelCase__       (      self       ,							a__       ,							a__ = -1       ,							a__ = -1       ,							a__ = False       ,							a__ = None       ,							)       ->       Mapping[str, Any]:

										'''simple docstring'''




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

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




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

										'''simple docstring'''




										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
										snake_case_			      =   tokenizer.num_special_tokens_to_add(a__							)
										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
										snake_case_			      =   [" ".join([tokenizer.unk_token]							) * seq_length] * batch_size
										snake_case_			      =   dict(tokenizer(a__       ,							return_tensors=a__							)							)
										return common_inputs




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

										'''simple docstring'''




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

										elif self.task == "causal-lm":
																	snake_case_			      =   self._generate_dummy_inputs_for_causal_lm(
																	    a__       ,							batch_size=a__       ,							seq_length=a__       ,							is_pair=a__       ,							framework=a__							)
										else:
																	snake_case_			      =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
																	    a__       ,							batch_size=a__       ,							seq_length=a__       ,							is_pair=a__       ,							framework=a__							)

										return common_inputs




			def      lowerCAmelCase__       (      self       ,							a__       ,							a__       ,							a__       ,							a__							)       ->       List[Any]:

										'''simple docstring'''




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






'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
    is_vision_available,
)


_SCREAMING_SNAKE_CASE							:     Dict   	=				{
    "configuration_mobilevitv2": [
        "MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "MobileViTV2Config",
        "MobileViTV2OnnxConfig",
    ],
}


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_SCREAMING_SNAKE_CASE							:     Optional[Any]   	=				[
		    "MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "MobileViTV2ForImageClassification",
		    "MobileViTV2ForSemanticSegmentation",
		    "MobileViTV2Model",
		    "MobileViTV2PreTrainedModel",
		]

if TYPE_CHECKING:
		from .configuration_mobilevitva import (
		    MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    MobileViTVaConfig,
		    MobileViTVaOnnxConfig,
		)

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_mobilevitva import (
				    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
				    MobileViTVaForImageClassification,
				    MobileViTVaForSemanticSegmentation,
				    MobileViTVaModel,
				    MobileViTVaPreTrainedModel,
				)

else:
		import sys

		_SCREAMING_SNAKE_CASE							:     Any   	=				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




