






"""simple docstring"""







from typing import Dict, Optional, Tuple

import torch






def 					_lowerCAmelCase				(					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	boundaries[1] - boundaries[0]
       UpperCAmelCase					=	boundaries + step / 2
       UpperCAmelCase					=	torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1						)]		,					dim=0						)
       return bin_centers






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					):







       '''simple docstring'''
       UpperCAmelCase					=	_calculate_bin_centers(lowerCAmelCase						)
       return (
           torch.sum(aligned_distance_error_probs * bin_centers		,					dim=-1						),
           bin_centers[-1],
       )






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase = 31		,					lowerCAmelCase = 64		,					**lowerCAmelCase		,					):







       '''simple docstring'''
       UpperCAmelCase					=	torch.linspace(0		,					lowerCAmelCase		,					steps=(no_bins - 1)		,					device=logits.device						)

       UpperCAmelCase					=	torch.nn.functional.softmax(lowerCAmelCase		,					dim=-1						)
       UpperCAmelCase			,				UpperCAmelCase					=	_calculate_expected_aligned_error(
           alignment_confidence_breaks=lowerCAmelCase		,					aligned_distance_error_probs=lowerCAmelCase		,					)

       return {
           "aligned_confidence_probs": aligned_confidence_probs,
           "predicted_aligned_error": predicted_aligned_error,
           "max_predicted_aligned_error": max_predicted_aligned_error,
       }






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase = None		,					lowerCAmelCase = 31		,					lowerCAmelCase = 64		,					lowerCAmelCase = 1e-8		,					**lowerCAmelCase		,					):







       '''simple docstring'''
       if residue_weights is None:
              UpperCAmelCase					=	logits.new_ones(logits.shape[-2]						)

       UpperCAmelCase					=	torch.linspace(0		,					lowerCAmelCase		,					steps=(no_bins - 1)		,					device=logits.device						)

       UpperCAmelCase					=	_calculate_bin_centers(lowerCAmelCase						)
       torch.sum(lowerCAmelCase						)
       UpperCAmelCase					=	logits.shape[-2]
       UpperCAmelCase					=	max(lowerCAmelCase		,					19						)

       UpperCAmelCase					=	1.24 * (clipped_n - 15) ** (1.0 / 3) - 1.8

       UpperCAmelCase					=	torch.nn.functional.softmax(lowerCAmelCase		,					dim=-1						)

       UpperCAmelCase					=	1.0 / (1 + (bin_centers**2) / (da**2))
       UpperCAmelCase					=	torch.sum(probs * tm_per_bin		,					dim=-1						)

       UpperCAmelCase					=	residue_weights / (eps + residue_weights.sum())
       UpperCAmelCase					=	torch.sum(predicted_tm_term * normed_residue_mask		,					dim=-1						)

       UpperCAmelCase					=	per_alignment * residue_weights

       UpperCAmelCase					=	(weighted == torch.max(lowerCAmelCase						)).nonzero()[0]
       return per_alignment[tuple(lowerCAmelCase						)]








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


lowercase     :							List[str]	     =			logging.get_logger(__name__)

lowercase     :							int	     =			{
    '''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 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ):


						'''simple docstring'''




						A   :      List[str]						 =    'bigbird_pegasus'
						A   :      Optional[int]						 =    ['past_key_values']
						A   :      Tuple						 =    {
						    'num_attention_heads': 'encoder_attention_heads',
						    'hidden_size': 'd_model',
						    'attention_probs_dropout_prob': 'attention_dropout',
						}






						def __init__(		self						, _SCREAMING_SNAKE_CASE=9_6103						, _SCREAMING_SNAKE_CASE=4096						, _SCREAMING_SNAKE_CASE=16						, _SCREAMING_SNAKE_CASE=4096						, _SCREAMING_SNAKE_CASE=16						, _SCREAMING_SNAKE_CASE=16						, _SCREAMING_SNAKE_CASE=4096						, _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_new"						, _SCREAMING_SNAKE_CASE=1024						, _SCREAMING_SNAKE_CASE=0.1						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=0.02						, _SCREAMING_SNAKE_CASE=2						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=True						, _SCREAMING_SNAKE_CASE=0						, _SCREAMING_SNAKE_CASE=2						, _SCREAMING_SNAKE_CASE=1						, _SCREAMING_SNAKE_CASE="block_sparse"						, _SCREAMING_SNAKE_CASE=64						, _SCREAMING_SNAKE_CASE=3						, _SCREAMING_SNAKE_CASE=False						, **_SCREAMING_SNAKE_CASE						, )       -> Dict:
								snake_case_ :							str         =				vocab_size
								snake_case_ :							Tuple         =				max_position_embeddings
								snake_case_ :							List[Any]         =				d_model
								snake_case_ :							List[str]         =				encoder_ffn_dim
								snake_case_ :							List[str]         =				encoder_layers
								snake_case_ :							int         =				encoder_attention_heads
								snake_case_ :							Tuple         =				decoder_ffn_dim
								snake_case_ :							List[Any]         =				decoder_layers
								snake_case_ :							Optional[int]         =				decoder_attention_heads
								snake_case_ :							Any         =				dropout
								snake_case_ :							Tuple         =				attention_dropout
								snake_case_ :							List[str]         =				activation_dropout
								snake_case_ :							Any         =				activation_function
								snake_case_ :							Any         =				init_std
								snake_case_ :							Optional[int]         =				encoder_layerdrop
								snake_case_ :							Optional[Any]         =				decoder_layerdrop
								snake_case_ :							Tuple         =				classifier_dropout
								snake_case_ :							str         =				use_cache
								snake_case_ :							Dict         =				encoder_layers
								snake_case_ :							List[str]         =				scale_embedding  # scale factor will be sqrt(d_model) if True

								# extra config
								snake_case_ :							Union[str, Any]         =				attention_type
								snake_case_ :							Tuple         =				block_size
								snake_case_ :							Any         =				num_random_blocks
								snake_case_ :							List[str]         =				use_bias

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




class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ):


						'''simple docstring'''




						@property
						def 					_lowerCAmelCase				(		self )       -> Mapping[str, Mapping[int, str]]:
								if self.task in ["default", "seq2seq-lm"]:
										snake_case_ :							str         =				OrderedDict(
										    [
										        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
										        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
										    ] )

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = None						, )       -> Mapping[str, Any]:
								snake_case_ :							List[str]         =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
								    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )

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

										snake_case_ :							Optional[Any]         =				torch.cat(
										    [common_inputs["decoder_attention_mask"], torch.ones(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )]						, dim=1 )

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

										snake_case_ :							str         =				common_inputs["attention_mask"].dtype
										snake_case_ :							Union[str, Any]         =				torch.cat(
										    [common_inputs["attention_mask"], torch.ones(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, dtype=_SCREAMING_SNAKE_CASE )]						, dim=1 )
										snake_case_ :							Optional[int]         =				[
										    (torch.zeros(_SCREAMING_SNAKE_CASE ), torch.zeros(_SCREAMING_SNAKE_CASE )) for _ in range(_SCREAMING_SNAKE_CASE )
										]
								return common_inputs






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = 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
								snake_case_ :							List[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
								snake_case_ :							Any         =				tokenizer.num_special_tokens_to_add(_SCREAMING_SNAKE_CASE )
								snake_case_ :							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
								snake_case_ :							List[Any]         =				[" ".join([tokenizer.unk_token] ) * seq_length] * batch_size
								snake_case_ :							Tuple         =				dict(tokenizer(_SCREAMING_SNAKE_CASE						, return_tensors=_SCREAMING_SNAKE_CASE ) )
								return common_inputs






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = None						, )       -> Mapping[str, Any]:
								if self.task in ["default", "seq2seq-lm"]:
										snake_case_ :							List[Any]         =				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 )

								elif self.task == "causal-lm":
										snake_case_ :							Optional[Any]         =				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 )
								else:
										snake_case_ :							Any         =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    _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 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Any:
								if self.task in ["default", "seq2seq-lm"]:
										snake_case_ :							int         =				super()._flatten_past_key_values_(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
								else:
										snake_case_ :							Optional[Any]         =				super(_SCREAMING_SNAKE_CASE						, self )._flatten_past_key_values_(
										    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
