






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

__a							:						Dict   							=			{
    '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 							_SCREAMING_SNAKE_CASE    (	__snake_case				):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE     ='bigbird_pegasus'
							_SCREAMING_SNAKE_CASE     =['past_key_values']
							_SCREAMING_SNAKE_CASE     ={
							    'num_attention_heads': 'encoder_attention_heads',
							    'hidden_size': 'd_model',
							    'attention_probs_dropout_prob': 'attention_dropout',
							}




							def __init__( self:  List[Any]					,		__A:  Any=96103					,		__A:  Optional[int]=4096					,		__A:  int=16					,		__A:  Any=4096					,		__A:  Any=16					,		__A:  Optional[int]=16					,		__A:  Optional[int]=4096					,		__A:  List[str]=16					,		__A:  str=0.0					,		__A:  Tuple=0.0					,		__A:  Optional[Any]=True					,		__A:  List[Any]=True					,		__A:  Optional[Any]="gelu_new"					,		__A:  Optional[int]=1024					,		__A:  Tuple=0.1					,		__A:  Tuple=0.0					,		__A:  int=0.0					,		__A:  List[Any]=0.0_2					,		__A:  Any=2					,		__A:  Union[str, Any]=0.0					,		__A:  Optional[int]=True					,		__A:  List[str]=0					,		__A:  Union[str, Any]=2					,		__A:  str=1					,		__A:  Optional[Any]="block_sparse"					,		__A:  str=64					,		__A:  Any=3					,		__A:  List[Any]=False					,		**__A:  Union[str, Any]					,		):

												'''simple docstring'''

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

												# extra config
												a__ 							=	attention_type
												a__ 							=	block_size
												a__ 							=	num_random_blocks
												a__ 							=	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 							_SCREAMING_SNAKE_CASE    (	__snake_case				):





							"""simple docstring"""





							@property
							def 			lowercase		( self:  Optional[int]     ):

												'''simple docstring'''

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

																	if self.use_past:
																						a__ 							=	{0: '''batch'''}
																						a__ 							=	{0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
																	else:
																						a__ 							=	{0: '''batch''', 1: '''decoder_sequence'''}
																						a__ 							=	{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.
																	a__ 							=	OrderedDict(
																	    [
																	        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
																	        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
																	    ]     )
																	if self.use_past:
																						a__							,a__ 							=	self.num_layers
																						for i in range(__A     ):
																											a__ 							=	{0: '''batch''', 2: '''past_sequence + sequence'''}
																											a__ 							=	{0: '''batch''', 2: '''past_sequence + sequence'''}
												else:
																	a__ 							=	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 			lowercase		( self:  Any     ):

												'''simple docstring'''

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




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

												'''simple docstring'''

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

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

																	a__ 							=	torch.cat(
																	    [common_inputs['''decoder_attention_mask'''], torch.ones(__A					,		__A     )]					,		dim=1     )

																	a__ 							=	[]
																	# If the number of encoder and decoder layers are present in the model configuration, both are considered
																	a__							,a__ 							=	self.num_layers
																	a__ 							=	min(__A					,		__A     )
																	a__ 							=	max(__A					,		__A     ) - min_num_layers
																	a__ 							=	'''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.
																	a__ 							=	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 			lowercase		( self:  List[Any]					,		__A:  PreTrainedTokenizer					,		__A:  int = -1					,		__A:  int = -1					,		__A:  bool = False					,		__A:  Optional[TensorType] = None					,		):

												'''simple docstring'''

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

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




							def 			lowercase		( self:  int					,		__A:  PreTrainedTokenizer					,		__A:  int = -1					,		__A:  int = -1					,		__A:  bool = False					,		__A:  Optional[TensorType] = None					,		):

												'''simple docstring'''

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




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

												'''simple docstring'''

												if self.task in ["default", "seq2seq-lm"]:
																	a__ 							=	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":
																	a__ 							=	self._generate_dummy_inputs_for_causal_lm(
																	    __A					,		batch_size=__A					,		seq_length=__A					,		is_pair=__A					,		framework=__A     )
												else:
																	a__ 							=	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 			lowercase		( self:  Optional[Any]					,		__A:  int					,		__A:  List[str]					,		__A:  Any					,		__A:  int     ):

												'''simple docstring'''

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










"""simple docstring"""







# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
	import tensorflow as tf







class 							_SCREAMING_SNAKE_CASE    (	TensorFormatter[Mapping, 'tf.Tensor', Mapping]				):





							"""simple docstring"""





							def __init__( self:  Union[str, Any]					,		__A:  Union[str, Any]=None					,		**__A:  Union[str, Any]     ):

												'''simple docstring'''

												super().__init__(features=__A     )
												a__ 							=	tf_tensor_kwargs
												import tensorflow as tf  # noqa: import tf at initialization




							def 			lowercase		( self:  Optional[int]					,		__A:  List[Any]     ):

												'''simple docstring'''

												import tensorflow as tf

												if isinstance(__A					,		__A     ) and column:
																	if all(
																	    isinstance(__A					,		tf.Tensor     ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column     ):
																						return tf.stack(__A     )
																	elif all(
																	    isinstance(__A					,		(tf.Tensor, tf.RaggedTensor)     ) and x.ndim == 1 and x.dtype == column[0].dtype
																	    for x in column     ):
																						# only rag 1-D tensors, otherwise some dimensions become ragged even though they were consolidated
																						return tf.ragged.stack(__A     )

												return column




							def 			lowercase		( self:  int					,		__A:  Optional[Any]     ):

												'''simple docstring'''

												import tensorflow as tf

												if value is None:
																	return value

												a__ 							=	{}

												if isinstance(__A					,		(np.number, np.ndarray)     ) and np.issubdtype(value.dtype					,		np.integer     ):
																	a__ 							=	{'''dtype''': tf.intaa}
												elif isinstance(__A					,		(np.number, np.ndarray)     ) and np.issubdtype(value.dtype					,		np.floating     ):
																	a__ 							=	{'''dtype''': tf.floataa}
												elif config.PIL_AVAILABLE and "PIL" in sys.modules:
																	import PIL.Image

																	if isinstance(__A					,		PIL.Image.Image     ):
																						a__ 							=	np.asarray(__A     )

												return tf.convert_to_tensor(__A					,		**{**default_dtype, **self.tf_tensor_kwargs}     )




							def 			lowercase		( self:  List[Any]					,		__A:  str     ):

												'''simple docstring'''

												import tensorflow as tf

												# support for torch, tf, jax etc.
												if config.TORCH_AVAILABLE and "torch" in sys.modules:
																	import torch

																	if isinstance(__A					,		torch.Tensor     ):
																						return self._tensorize(data_struct.detach().cpu().numpy()[()]     )
												if hasattr(__A					,		'''__array__'''     ) and not isinstance(__A					,		tf.Tensor     ):
																	a__ 							=	data_struct.__array__()
												# support for nested types like struct of list of struct
												if isinstance(__A					,		np.ndarray     ):
																	if data_struct.dtype == object:  # tf tensors cannot be instantied from an array of objects
																						return self._consolidate([self.recursive_tensorize(__A     ) for substruct in data_struct]     )
												elif isinstance(__A					,		(list, tuple)     ):
																	return self._consolidate([self.recursive_tensorize(__A     ) for substruct in data_struct]     )
												return self._tensorize(__A     )




							def 			lowercase		( self:  int					,		__A:  dict     ):

												'''simple docstring'''

												return map_nested(self._recursive_tensorize					,		__A					,		map_list=__A     )




							def 			lowercase		( self:  Tuple					,		__A:  pa.Table     ):

												'''simple docstring'''

												a__ 							=	self.numpy_arrow_extractor().extract_row(__A     )
												a__ 							=	self.python_features_decoder.decode_row(__A     )
												return self.recursive_tensorize(__A     )




							def 			lowercase		( self:  Dict					,		__A:  pa.Table     ):

												'''simple docstring'''

												a__ 							=	self.numpy_arrow_extractor().extract_column(__A     )
												a__ 							=	self.python_features_decoder.decode_column(__A					,		pa_table.column_names[0]     )
												a__ 							=	self.recursive_tensorize(__A     )
												a__ 							=	self._consolidate(__A     )
												return column





							def 			lowercase		( self:  Tuple					,		__A:  pa.Table     ):

												'''simple docstring'''

												a__ 							=	self.numpy_arrow_extractor().extract_batch(__A     )
												a__ 							=	self.python_features_decoder.decode_batch(__A     )
												a__ 							=	self.recursive_tensorize(__A     )
												for column_name in batch:
																	a__ 							=	self._consolidate(batch[column_name]     )
												return batch


