


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


__magic_name__	:   Tuple          =							logging.get_logger(__name__)

__magic_name__	:   List[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      (lowerCamelCase     ):
							__a				   =					'''bigbird_pegasus'''
							__a				   =					['''past_key_values''']
							__a				   =					{
							    '''num_attention_heads''': '''encoder_attention_heads''',
							    '''hidden_size''': '''d_model''',
							    '''attention_probs_dropout_prob''': '''attention_dropout''',
							}




							def __init__(    self:						Any    ,       A_:						Optional[Any]=9_61_03    ,       A_:						Optional[Any]=40_96    ,       A_:						List[str]=16    ,       A_:						Optional[Any]=40_96    ,       A_:						Dict=16    ,       A_:						Any=16    ,       A_:						Dict=40_96    ,       A_:						Optional[int]=16    ,       A_:						Any=0.0    ,       A_:						List[str]=0.0    ,       A_:						Dict=True    ,       A_:						int=True    ,       A_:						Union[str, Any]="gelu_new"    ,       A_:						Tuple=10_24    ,       A_:						Optional[int]=0.1    ,       A_:						Dict=0.0    ,       A_:						int=0.0    ,       A_:						str=0.02    ,       A_:						Optional[int]=2    ,       A_:						List[str]=0.0    ,       A_:						Any=True    ,       A_:						Any=0    ,       A_:						Tuple=2    ,       A_:						int=1    ,       A_:						Optional[Any]="block_sparse"    ,       A_:						Any=64    ,       A_:						Any=3    ,       A_:						int=False    ,       **A_:						Optional[int]    ,       ):
										__lowerCamelCase        =			vocab_size
										__lowerCamelCase        =			max_position_embeddings
										__lowerCamelCase        =			d_model
										__lowerCamelCase        =			encoder_ffn_dim
										__lowerCamelCase        =			encoder_layers
										__lowerCamelCase        =			encoder_attention_heads
										__lowerCamelCase        =			decoder_ffn_dim
										__lowerCamelCase        =			decoder_layers
										__lowerCamelCase        =			decoder_attention_heads
										__lowerCamelCase        =			dropout
										__lowerCamelCase        =			attention_dropout
										__lowerCamelCase        =			activation_dropout
										__lowerCamelCase        =			activation_function
										__lowerCamelCase        =			init_std
										__lowerCamelCase        =			encoder_layerdrop
										__lowerCamelCase        =			decoder_layerdrop
										__lowerCamelCase        =			classifier_dropout
										__lowerCamelCase        =			use_cache
										__lowerCamelCase        =			encoder_layers
										__lowerCamelCase        =			scale_embedding  # scale factor will be sqrt(d_model) if True

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




							@property
							def      __a      (    self:						List[Any]      ):
										if self.task in ["default", "seq2seq-lm"]:
													__lowerCamelCase        =			OrderedDict(
													    [
													        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
													        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
													    ]      )

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




							def      __a      (    self:						Dict    ,       A_:						PreTrainedTokenizer    ,       A_:						int = -1    ,       A_:						int = -1    ,       A_:						bool = False    ,       A_:						Optional[TensorType] = None    ,       ):
										__lowerCamelCase        =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    A_    ,       A_    ,       A_    ,       A_    ,       A_      )

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

													__lowerCamelCase        =			torch.cat(
													    [common_inputs["""decoder_attention_mask"""], torch.ones(A_    ,       A_      )]    ,       dim=1      )

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

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




							def      __a      (    self:						List[Any]    ,       A_:						PreTrainedTokenizer    ,       A_:						int = -1    ,       A_:						int = -1    ,       A_:						bool = False    ,       A_:						Optional[TensorType] = None    ,       ):
										# 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
										__lowerCamelCase        =			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
										__lowerCamelCase        =			tokenizer.num_special_tokens_to_add(A_      )
										__lowerCamelCase        =			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
										__lowerCamelCase        =			[""" """.join([tokenizer.unk_token]      ) * seq_length] * batch_size
										__lowerCamelCase        =			dict(tokenizer(A_    ,       return_tensors=A_      )      )
										return common_inputs




							def      __a      (    self:						Dict    ,       A_:						PreTrainedTokenizer    ,       A_:						int = -1    ,       A_:						int = -1    ,       A_:						bool = False    ,       A_:						Optional[TensorType] = None    ,       ):
										if self.task in ["default", "seq2seq-lm"]:
													__lowerCamelCase        =			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":
													__lowerCamelCase        =			self._generate_dummy_inputs_for_causal_lm(
													    A_    ,       batch_size=A_    ,       seq_length=A_    ,       is_pair=A_    ,       framework=A_      )
										else:
													__lowerCamelCase        =			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      __a      (    self:						str    ,       A_:						Optional[Any]    ,       A_:						str    ,       A_:						Optional[int]    ,       A_:						Optional[Any]      ):
										if self.task in ["default", "seq2seq-lm"]:
													__lowerCamelCase        =			super()._flatten_past_key_values_(A_    ,       A_    ,       A_    ,       A_      )
										else:
													__lowerCamelCase        =			super(A_    ,       self      )._flatten_past_key_values_(
													    A_    ,       A_    ,       A_    ,       A_      )








"""simple docstring"""




import unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types




class  __snake_case      (unittest.TestCase ,		lowerCamelCase     ):




							def      __a      (    self:						Optional[Any]      ):
										__lowerCamelCase        =			load_tool("""translation"""      )
										self.tool.setup()
										__lowerCamelCase        =			load_tool("""translation"""    ,       remote=A_      )




							def      __a      (    self:						int      ):
										__lowerCamelCase        =			self.tool("""Hey, what's up?"""    ,       src_lang="""English"""    ,       tgt_lang="""French"""      )
										self.assertEqual(A_    ,       """- Hé, comment ça va?"""      )




							def      __a      (    self:						int      ):
										__lowerCamelCase        =			self.remote_tool("""Hey, what's up?"""    ,       src_lang="""English"""    ,       tgt_lang="""French"""      )
										self.assertEqual(A_    ,       """- Hé, comment ça va?"""      )




							def      __a      (    self:						Union[str, Any]      ):
										__lowerCamelCase        =			self.tool(text="""Hey, what's up?"""    ,       src_lang="""English"""    ,       tgt_lang="""French"""      )
										self.assertEqual(A_    ,       """- Hé, comment ça va?"""      )




							def      __a      (    self:						Any      ):
										__lowerCamelCase        =			self.remote_tool(text="""Hey, what's up?"""    ,       src_lang="""English"""    ,       tgt_lang="""French"""      )
										self.assertEqual(A_    ,       """- Hé, comment ça va?"""      )




							def      __a      (    self:						List[Any]      ):
										__lowerCamelCase        =			["""Hey, what's up?""", """English""", """Spanish"""]
										__lowerCamelCase        =			self.tool(*A_      )

										# There is a single output
										if len(self.tool.outputs      ) == 1:
													__lowerCamelCase        =			[outputs]

										self.assertListEqual(output_types(A_      )    ,       self.tool.outputs      )




							def      __a      (    self:						Union[str, Any]      ):
										__lowerCamelCase        =			["""Hey, what's up?""", """English""", """Spanish"""]
										__lowerCamelCase        =			self.tool(*A_      )

										if not isinstance(A_    ,       A_      ):
													__lowerCamelCase        =			[outputs]

										self.assertEqual(len(A_      )    ,       len(self.tool.outputs      )      )

										for output, output_type in zip(A_    ,       self.tool.outputs      ):
													__lowerCamelCase        =			AGENT_TYPE_MAPPING[output_type]
													self.assertTrue(isinstance(A_    ,       A_      )      )




							def      __a      (    self:						Dict      ):
										__lowerCamelCase        =			["""Hey, what's up?""", """English""", """Spanish"""]

										__lowerCamelCase        =			[]

										for _input, input_type in zip(A_    ,       self.tool.inputs      ):
													if isinstance(A_    ,       A_      ):
																_inputs.append([AGENT_TYPE_MAPPING[_input_type](_input      ) for _input_type in input_type]      )
													else:
																_inputs.append(AGENT_TYPE_MAPPING[input_type](_input      )      )

        # Should not raise an error
										__lowerCamelCase        =			self.tool(*A_      )

										if not isinstance(A_    ,       A_      ):
													__lowerCamelCase        =			[outputs]

										self.assertEqual(len(A_      )    ,       len(self.tool.outputs      )      )




