



"""simple docstring"""






import random
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
				import torch

if is_vision_available():
				import PIL

				from transformers import FlavaImageProcessor
				from transformers.image_utils import PILImageResampling
				from transformers.models.flava.image_processing_flava import (
				    FLAVA_CODEBOOK_MEAN,
				    FLAVA_CODEBOOK_STD,
				    FLAVA_IMAGE_MEAN,
				    FLAVA_IMAGE_STD,
				)
else:
				__UpperCAmelCase      :      str   					=					None

class 	__lowerCAmelCase							(unittest.TestCase  ):


					'''simple docstring'''


					def __init__(      self    ,       a    ,       a=7    ,       a=3    ,       a=30    ,       a=4_00    ,       a=True    ,       a=None    ,       a=True    ,       a=None    ,       a=None    ,       a=True    ,       a=1 / 2_55    ,       a=True    ,       a=FLAVA_IMAGE_MEAN    ,       a=FLAVA_IMAGE_STD    ,       a=14    ,       a=75    ,       a=None    ,       a=16    ,       a=0.3    ,       a=None    ,       a=True    ,       a=None    ,       a=None    ,       a=True    ,       a=None    ,       a=True    ,       a=True    ,       a=FLAVA_CODEBOOK_MEAN    ,       a=FLAVA_CODEBOOK_STD    ,       ):




										"""simple docstring"""





										snake_case_						:List[Any]								=						size if size is not None else {"height": 2_24, "width": 2_24}
										snake_case_						:List[Any]								=						crop_size if crop_size is not None else {"height": 2_24, "width": 2_24}
										snake_case_						:Tuple								=						codebook_size if codebook_size is not None else {"height": 1_12, "width": 1_12}
										snake_case_						:str								=						codebook_crop_size if codebook_crop_size is not None else {"height": 1_12, "width": 1_12}

										snake_case_						:List[Any]								=						parent
										snake_case_						:Dict								=						batch_size
										snake_case_						:Optional[int]								=						num_channels
										snake_case_						:int								=						do_resize
										snake_case_						:Union[str, Any]								=						do_rescale
										snake_case_						:str								=						rescale_factor
										snake_case_						:Optional[Any]								=						min_resolution
										snake_case_						:Optional[Any]								=						max_resolution
										snake_case_						:Union[str, Any]								=						size
										snake_case_						:Optional[int]								=						resample if resample is not None else PILImageResampling.BICUBIC
										snake_case_						:str								=						do_normalize
										snake_case_						:Tuple								=						image_mean
										snake_case_						:Tuple								=						image_std
										snake_case_						:Tuple								=						do_center_crop
										snake_case_						:List[str]								=						crop_size

										snake_case_						:Union[str, Any]								=						input_size_patches
										snake_case_						:Any								=						total_mask_patches
										snake_case_						:Optional[int]								=						mask_group_max_patches
										snake_case_						:Any								=						mask_group_min_patches
										snake_case_						:Union[str, Any]								=						mask_group_min_aspect_ratio
										snake_case_						:str								=						mask_group_max_aspect_ratio

										snake_case_						:List[str]								=						codebook_do_resize
										snake_case_						:int								=						codebook_size
										snake_case_						:Optional[int]								=						codebook_resample if codebook_resample is not None else PILImageResampling.LANCZOS
										snake_case_						:Any								=						codebook_do_center_crop
										snake_case_						:Dict								=						codebook_crop_size
										snake_case_						:List[str]								=						codebook_do_map_pixels
										snake_case_						:Any								=						codebook_do_normalize
										snake_case_						:Optional[int]								=						codebook_image_mean
										snake_case_						:Dict								=						codebook_image_std







					def 		_a (      self   ):




										"""simple docstring"""





										return {
										    "image_mean": self.image_mean,
										    "image_std": self.image_std,
										    "do_normalize": self.do_normalize,
										    "do_resize": self.do_resize,
										    "size": self.size,
										    "resample": self.resample,
										    "do_rescale": self.do_rescale,
										    "rescale_factor": self.rescale_factor,
										    "do_center_crop": self.do_center_crop,
										    "crop_size": self.crop_size,
										    "input_size_patches": self.input_size_patches,
										    "total_mask_patches": self.total_mask_patches,
										    "mask_group_max_patches": self.mask_group_max_patches,
										    "mask_group_min_patches": self.mask_group_min_patches,
										    "mask_group_min_aspect_ratio": self.mask_group_min_aspect_ratio,
										    "mask_group_max_aspect_ratio": self.mask_group_min_aspect_ratio,
										    "codebook_do_resize": self.codebook_do_resize,
										    "codebook_size": self.codebook_size,
										    "codebook_resample": self.codebook_resample,
										    "codebook_do_center_crop": self.codebook_do_center_crop,
										    "codebook_crop_size": self.codebook_crop_size,
										    "codebook_do_map_pixels": self.codebook_do_map_pixels,
										    "codebook_do_normalize": self.codebook_do_normalize,
										    "codebook_image_mean": self.codebook_image_mean,
										    "codebook_image_std": self.codebook_image_std,
										}







					def 		_a (      self   ):




										"""simple docstring"""





										return (self.size["height"], self.size["width"])







					def 		_a (      self   ):




										"""simple docstring"""





										return (
										    (self.input_size_patches, self.input_size_patches)
										    if not isinstance(self.input_size_patches    ,       a   )
										    else self.input_size_patches
										)







					def 		_a (      self   ):




										"""simple docstring"""





										return (self.codebook_size["height"], self.codebook_size["width"])



@require_torch
@require_vision
class 	__lowerCAmelCase							(__UpperCamelCase					,unittest.TestCase  ):


					'''simple docstring'''


					a__        =     FlavaImageProcessor if is_vision_available() else None
					a__        =     None







					def 		_a (      self   ):




										"""simple docstring"""





										snake_case_						:int								=						FlavaImageProcessingTester(self   )







					@property
					def 		_a (      self   ):




										"""simple docstring"""





										return self.image_processor_tester.prepare_image_processor_dict()







					def 		_a (      self   ):




										"""simple docstring"""





										snake_case_						:List[str]								=						self.image_processing_class(**self.image_processor_dict   )
										self.assertTrue(hasattr(a    ,       "image_mean"   )   )
										self.assertTrue(hasattr(a    ,       "image_std"   )   )
										self.assertTrue(hasattr(a    ,       "do_normalize"   )   )
										self.assertTrue(hasattr(a    ,       "do_resize"   )   )
										self.assertTrue(hasattr(a    ,       "resample"   )   )
										self.assertTrue(hasattr(a    ,       "crop_size"   )   )
										self.assertTrue(hasattr(a    ,       "do_center_crop"   )   )
										self.assertTrue(hasattr(a    ,       "do_rescale"   )   )
										self.assertTrue(hasattr(a    ,       "rescale_factor"   )   )
										self.assertTrue(hasattr(a    ,       "masking_generator"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_do_resize"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_size"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_resample"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_do_center_crop"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_crop_size"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_do_map_pixels"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_do_normalize"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_image_mean"   )   )
										self.assertTrue(hasattr(a    ,       "codebook_image_std"   )   )







					def 		_a (      self   ):




										"""simple docstring"""





										snake_case_						:Tuple								=						self.image_processing_class.from_dict(self.image_processor_dict   )
										self.assertEqual(image_processor.size    ,       {"height": 2_24, "width": 2_24}   )
										self.assertEqual(image_processor.crop_size    ,       {"height": 2_24, "width": 2_24}   )
										self.assertEqual(image_processor.codebook_size    ,       {"height": 1_12, "width": 1_12}   )
										self.assertEqual(image_processor.codebook_crop_size    ,       {"height": 1_12, "width": 1_12}   )

										snake_case_						:Dict								=						self.image_processing_class.from_dict(
										    self.image_processor_dict    ,       size=42    ,       crop_size=84    ,       codebook_size=33    ,       codebook_crop_size=66   )
										self.assertEqual(image_processor.size    ,       {"height": 42, "width": 42}   )
										self.assertEqual(image_processor.crop_size    ,       {"height": 84, "width": 84}   )
										self.assertEqual(image_processor.codebook_size    ,       {"height": 33, "width": 33}   )
										self.assertEqual(image_processor.codebook_crop_size    ,       {"height": 66, "width": 66}   )







					def 		_a (      self   ):




										"""simple docstring"""





										pass







					def 		_a (      self   ):




										"""simple docstring"""





										snake_case_						:str								=						self.image_processing_class(**self.image_processor_dict   )
										# create random PIL images
										snake_case_						:Any								=						prepare_image_inputs(self.image_processor_tester    ,       equal_resolution=a   )
										for image in image_inputs:
															self.assertIsInstance(a    ,       PIL.Image.Image   )

										# Test not batched input
										snake_case_						:Dict								=						image_processing(image_inputs[0]    ,       return_tensors="pt"   )

										# Test no bool masked pos
										self.assertFalse("bool_masked_pos" in encoded_images   )

										snake_case_	,  snake_case_						:List[Any]								=						self.image_processor_tester.get_expected_image_size()

										self.assertEqual(
										    encoded_images.pixel_values.shape    ,       (1, self.image_processor_tester.num_channels, expected_height, expected_width)    ,       )

										# Test batched
										snake_case_						:List[str]								=						image_processing(a    ,       return_tensors="pt"   )
										snake_case_	,  snake_case_						:Optional[Any]								=						self.image_processor_tester.get_expected_image_size()

										# Test no bool masked pos
										self.assertFalse("bool_masked_pos" in encoded_images   )

										self.assertEqual(
										    encoded_images.pixel_values.shape    ,       (
										        self.image_processor_tester.batch_size,
										        self.image_processor_tester.num_channels,
										        expected_height,
										        expected_width,
										    )    ,       )







					def 		_a (      self    ,       a    ,       a   ):




										"""simple docstring"""





										snake_case_						:Any								=						self.image_processing_class(**self.image_processor_dict   )
										# create random tensors
										snake_case_						:Optional[int]								=						prepare_image_inputs(self.image_processor_tester    ,       equal_resolution=a    ,       **a   )
										for image in image_inputs:
															self.assertIsInstance(a    ,       a   )

										# Test not batched input
										snake_case_						:Tuple								=						image_processing(image_inputs[0]    ,       return_tensors="pt"   )

										snake_case_	,  snake_case_						:List[Any]								=						self.image_processor_tester.get_expected_image_size()
										self.assertEqual(
										    encoded_images.pixel_values.shape    ,       (1, self.image_processor_tester.num_channels, expected_height, expected_width)    ,       )

										snake_case_						:Optional[int]								=						image_processing(a    ,       return_image_mask=a    ,       return_tensors="pt"   )

										snake_case_	,  snake_case_						:Tuple								=						self.image_processor_tester.get_expected_image_size()
										self.assertEqual(
										    encoded_images.pixel_values.shape    ,       (
										        self.image_processor_tester.batch_size,
										        self.image_processor_tester.num_channels,
										        expected_height,
										        expected_width,
										    )    ,       )

										snake_case_	,  snake_case_						:Optional[int]								=						self.image_processor_tester.get_expected_mask_size()
										self.assertEqual(
										    encoded_images.bool_masked_pos.shape    ,       (
										        self.image_processor_tester.batch_size,
										        expected_height,
										        expected_width,
										    )    ,       )

										# Test batched
										snake_case_						:Optional[Any]								=						image_processing(a    ,       return_tensors="pt"   ).pixel_values

										snake_case_	,  snake_case_						:int								=						self.image_processor_tester.get_expected_image_size()
										self.assertEqual(
										    encoded_images.shape    ,       (
										        self.image_processor_tester.batch_size,
										        self.image_processor_tester.num_channels,
										        expected_height,
										        expected_width,
										    )    ,       )

										# Test masking
										snake_case_						:Optional[int]								=						image_processing(a    ,       return_image_mask=a    ,       return_tensors="pt"   )

										snake_case_	,  snake_case_						:List[Any]								=						self.image_processor_tester.get_expected_image_size()
										self.assertEqual(
										    encoded_images.pixel_values.shape    ,       (
										        self.image_processor_tester.batch_size,
										        self.image_processor_tester.num_channels,
										        expected_height,
										        expected_width,
										    )    ,       )

										snake_case_	,  snake_case_						:Tuple								=						self.image_processor_tester.get_expected_mask_size()
										self.assertEqual(
										    encoded_images.bool_masked_pos.shape    ,       (
										        self.image_processor_tester.batch_size,
										        expected_height,
										        expected_width,
										    )    ,       )







					def 		_a (      self   ):




										"""simple docstring"""





										self._test_call_framework(np.ndarray    ,       prepare_kwargs={"numpify": True}   )







					def 		_a (      self   ):




										"""simple docstring"""





										self._test_call_framework(torch.Tensor    ,       prepare_kwargs={"torchify": True}   )







					def 		_a (      self   ):




										"""simple docstring"""





										random.seed(12_34   )
										snake_case_						:List[str]								=						self.image_processing_class(**self.image_processor_dict   )
										snake_case_						:Any								=						prepare_image_inputs(self.image_processor_tester    ,       equal_resolution=a    ,       torchify=a   )

										# Test not batched input
										snake_case_						:Optional[int]								=						image_processing(image_inputs[0]    ,       return_image_mask=a    ,       return_tensors="pt"   )
										self.assertEqual(encoded_images.bool_masked_pos.sum().item()    ,       75   )







					def 		_a (      self   ):




										"""simple docstring"""





										snake_case_						:Optional[int]								=						self.image_processing_class(**self.image_processor_dict   )
										# create random PIL images
										snake_case_						:Optional[int]								=						prepare_image_inputs(self.image_processor_tester    ,       equal_resolution=a   )
										for image in image_inputs:
															self.assertIsInstance(a    ,       PIL.Image.Image   )

										# Test not batched input
										snake_case_						:Dict								=						image_processing(image_inputs[0]    ,       return_codebook_pixels=a    ,       return_tensors="pt"   )
										snake_case_	,  snake_case_						:str								=						self.image_processor_tester.get_expected_codebook_image_size()
										self.assertEqual(
										    encoded_images.codebook_pixel_values.shape    ,       (1, self.image_processor_tester.num_channels, expected_height, expected_width)    ,       )

										# Test batched
										snake_case_						:List[Any]								=						image_processing(a    ,       return_codebook_pixels=a    ,       return_tensors="pt"   )
										snake_case_	,  snake_case_						:Any								=						self.image_processor_tester.get_expected_codebook_image_size()
										self.assertEqual(
										    encoded_images.codebook_pixel_values.shape    ,       (
										        self.image_processor_tester.batch_size,
										        self.image_processor_tester.num_channels,
										        expected_height,
										        expected_width,
										    )    ,       )






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


__UpperCAmelCase      :      Union[str, Any]   					=					logging.get_logger(__name__)

__UpperCAmelCase      :      str   					=					{
    '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 	__lowerCAmelCase							(__UpperCamelCase  ):


					'''simple docstring'''


					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    ,       a=9_61_03    ,       a=40_96    ,       a=16    ,       a=40_96    ,       a=16    ,       a=16    ,       a=40_96    ,       a=16    ,       a=0.0    ,       a=0.0    ,       a=True    ,       a=True    ,       a="gelu_new"    ,       a=10_24    ,       a=0.1    ,       a=0.0    ,       a=0.0    ,       a=0.02    ,       a=2    ,       a=0.0    ,       a=True    ,       a=0    ,       a=2    ,       a=1    ,       a="block_sparse"    ,       a=64    ,       a=3    ,       a=False    ,       **a    ,       ):




										"""simple docstring"""





										snake_case_						:Optional[int]								=						vocab_size
										snake_case_						:Any								=						max_position_embeddings
										snake_case_						:List[Any]								=						d_model
										snake_case_						:Any								=						encoder_ffn_dim
										snake_case_						:str								=						encoder_layers
										snake_case_						:Union[str, Any]								=						encoder_attention_heads
										snake_case_						:Any								=						decoder_ffn_dim
										snake_case_						:Dict								=						decoder_layers
										snake_case_						:List[Any]								=						decoder_attention_heads
										snake_case_						:Any								=						dropout
										snake_case_						:Union[str, Any]								=						attention_dropout
										snake_case_						:Optional[int]								=						activation_dropout
										snake_case_						:int								=						activation_function
										snake_case_						:int								=						init_std
										snake_case_						:Union[str, Any]								=						encoder_layerdrop
										snake_case_						:Dict								=						decoder_layerdrop
										snake_case_						:Any								=						classifier_dropout
										snake_case_						:Any								=						use_cache
										snake_case_						:Any								=						encoder_layers
										snake_case_						:int								=						scale_embedding  # scale factor will be sqrt(d_model) if True

										# extra config
										snake_case_						:Optional[int]								=						attention_type
										snake_case_						:int								=						block_size
										snake_case_						:Optional[int]								=						num_random_blocks
										snake_case_						:List[str]								=						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 	__lowerCAmelCase							(__UpperCamelCase  ):


					'''simple docstring'''


					@property
					def 		_a (      self   ):




										"""simple docstring"""





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

															if self.use_past:
																				snake_case_						:int								=						{0: "batch"}
																				snake_case_						:List[str]								=						{0: "batch", 1: "past_decoder_sequence + sequence"}
															else:
																				snake_case_						:Optional[int]								=						{0: "batch", 1: "decoder_sequence"}
																				snake_case_						:Optional[Any]								=						{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_						:List[str]								=						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(a   ):
																									snake_case_						:int								=						{0: "batch", 2: "past_sequence + sequence"}
																									snake_case_						:List[Any]								=						{0: "batch", 2: "past_sequence + sequence"}
										else:
															snake_case_						:List[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 		_a (      self   ):




										"""simple docstring"""





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







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:int								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a    ,       a    ,       a    ,       a    ,       a   )

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

															snake_case_						:Union[str, Any]								=						torch.cat(
															    [common_inputs["decoder_attention_mask"], torch.ones(a    ,       a   )]    ,       dim=1   )

															snake_case_						:Union[str, Any]								=						[]
															# If the number of encoder and decoder layers are present in the model configuration, both are considered
															snake_case_	,  snake_case_						:Tuple								=						self.num_layers
															snake_case_						:int								=						min(a    ,       a   )
															snake_case_						:Optional[Any]								=						max(a    ,       a   ) - min_num_layers
															snake_case_						:List[Any]								=						"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_						:Any								=						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    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





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

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







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:Optional[int]								=						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_						:List[Any]								=						tokenizer.num_special_tokens_to_add(a   )
										snake_case_						:Tuple								=						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_						:Any								=						[" ".join([tokenizer.unk_token]   ) * seq_length] * batch_size
										snake_case_						:List[str]								=						dict(tokenizer(a    ,       return_tensors=a   )   )
										return common_inputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:Dict								=						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_						:Optional[Any]								=						self._generate_dummy_inputs_for_causal_lm(
															    a    ,       batch_size=a    ,       seq_length=a    ,       is_pair=a    ,       framework=a   )
										else:
															snake_case_						:List[Any]								=						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    ,       a    ,       a    ,       a    ,       a   ):




										"""simple docstring"""





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

