




'''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	):
       snake_case_					:	Tuple       =   """bigbird_pegasus"""
       snake_case_					:	List[Any]       =   ["""past_key_values"""]
       snake_case_					:	Union[str, Any]       =   {
           """num_attention_heads""": """encoder_attention_heads""",
           """hidden_size""": """d_model""",
           """attention_probs_dropout_prob""": """attention_dropout""",
       }

       def __init__(       self   ,						lowerCamelCase__=9_6103   ,						lowerCamelCase__=4096   ,						lowerCamelCase__=16   ,						lowerCamelCase__=4096   ,						lowerCamelCase__=16   ,						lowerCamelCase__=16   ,						lowerCamelCase__=4096   ,						lowerCamelCase__=16   ,						lowerCamelCase__=0.0   ,						lowerCamelCase__=0.0   ,						lowerCamelCase__=True   ,						lowerCamelCase__=True   ,						lowerCamelCase__="gelu_new"   ,						lowerCamelCase__=1024   ,						lowerCamelCase__=0.1   ,						lowerCamelCase__=0.0   ,						lowerCamelCase__=0.0   ,						lowerCamelCase__=0.02   ,						lowerCamelCase__=2   ,						lowerCamelCase__=0.0   ,						lowerCamelCase__=True   ,						lowerCamelCase__=0   ,						lowerCamelCase__=2   ,						lowerCamelCase__=1   ,						lowerCamelCase__="block_sparse"   ,						lowerCamelCase__=64   ,						lowerCamelCase__=3   ,						lowerCamelCase__=False   ,						**lowerCamelCase__   ,						)       ->   Any:

           """simple docstring"""



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

           # extra config
           __lowercase     							=		attention_type
           __lowercase     							=		block_size
           __lowercase     							=		num_random_blocks
           __lowercase     							=		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	):

       @property
       def        snake_case__     (       self	)       ->   Dict:

           """simple docstring"""



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

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

           """simple docstring"""



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

       def        snake_case__     (       self   ,						lowerCamelCase__   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = False   ,						lowerCamelCase__ = None   ,						)       ->   str:

           """simple docstring"""



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

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

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

               __lowercase     							=		[]
               # If the number of encoder and decoder layers are present in the model configuration, both are considered
               __lowercase							,__lowercase     							=		self.num_layers
               __lowercase     							=		min(__A   ,						__A	)
               __lowercase     							=		max(__A   ,						__A	) - min_num_layers
               __lowercase     							=		"""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.
               __lowercase     							=		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   ,						lowerCamelCase__   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = False   ,						lowerCamelCase__ = None   ,						)       ->   int:

           """simple docstring"""



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

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

       def        snake_case__     (       self   ,						lowerCamelCase__   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = False   ,						lowerCamelCase__ = None   ,						)       ->   Union[str, Any]:

           """simple docstring"""



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

       def        snake_case__     (       self   ,						lowerCamelCase__   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = -1   ,						lowerCamelCase__ = False   ,						lowerCamelCase__ = None   ,						)       ->   Union[str, Any]:

           """simple docstring"""



           if self.task in ["default", "seq2seq-lm"]:
               __lowercase     							=		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":
               __lowercase     							=		self._generate_dummy_inputs_for_causal_lm(
                   __A   ,						batch_size=__A   ,						seq_length=__A   ,						is_pair=__A   ,						framework=__A	)
           else:
               __lowercase     							=		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        snake_case__     (       self   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__	)       ->   Dict:

           """simple docstring"""



           if self.task in ["default", "seq2seq-lm"]:
               __lowercase     							=		super()._flatten_past_key_values_(__A   ,						__A   ,						__A   ,						__A	)
           else:
               __lowercase     							=		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


