




'''simple docstring'''







import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class 					lowerCamelCase_							(	unittest.TestCase							):




					'''simple docstring'''






					def     __SCREAMING_SNAKE_CASE     (   self			):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	1
							_snake_case     	=	3
							_snake_case     	=	(32, 32)

							_snake_case     	=	floats_tensor((batch_size, num_channels) + sizes,  rng=random.Random(0			)			).to(__UpperCamelCase			)
							return image






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							torch.manual_seed(0			)
							_snake_case     	=	UNetaDConditionModel(
							    block_out_channels=(32, 64),  layers_per_block=2,  sample_size=32,  in_channels=4,  out_channels=4,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),  up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),  cross_attention_dim=32,  )
							return model






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							torch.manual_seed(0			)
							_snake_case     	=	AutoencoderKL(
							    block_out_channels=[32, 64],  in_channels=3,  out_channels=3,  down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],  up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],  latent_channels=4,  )
							return model






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							torch.manual_seed(0			)
							_snake_case     	=	CLIPTextConfig(
							    bos_token_id=0,  eos_token_id=2,  hidden_size=32,  intermediate_size=37,  layer_norm_eps=1E-05,  num_attention_heads=4,  num_hidden_layers=5,  pad_token_id=1,  vocab_size=1_000,  )
							return CLIPTextModel(__UpperCamelCase			)






					@property
					def     __SCREAMING_SNAKE_CASE     (   self			):
							def extract(*__UpperCamelCase,  **__UpperCamelCase			):





									class 					lowerCamelCase_							:




														'''simple docstring'''






														def __init__(   self			):
																_snake_case     	=	torch.ones([0]			)






														def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase			):
																self.pixel_values.to(__UpperCamelCase			)
																return self




									return Out()

							return extract






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	"cpu"  # ensure determinism for the device-dependent torch.Generator
							_snake_case     	=	self.dummy_cond_unet
							_snake_case     	=	DDIMScheduler(
							    beta_start=0.0_00_85,  beta_end=0.0_12,  beta_schedule="scaled_linear",  clip_sample=__UpperCamelCase,  set_alpha_to_one=__UpperCamelCase,  )

							_snake_case     	=	self.dummy_vae
							_snake_case     	=	self.dummy_text_encoder
							_snake_case     	=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

							# make sure here that pndm scheduler skips prk
							_snake_case     	=	StableDiffusionPipeline(
							    unet=__UpperCamelCase,  scheduler=__UpperCamelCase,  vae=__UpperCamelCase,  text_encoder=__UpperCamelCase,  tokenizer=__UpperCamelCase,  safety_checker=__UpperCamelCase,  feature_extractor=self.dummy_extractor,  )
							_snake_case     	=	sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							_snake_case     	=	"A painting of a squirrel eating a burger"

							_snake_case     	=	torch.Generator(device=__UpperCamelCase			).manual_seed(0			)
							_snake_case     	=	sd_pipe([prompt],  generator=__UpperCamelCase,  guidance_scale=6.0,  num_inference_steps=2,  output_type="np"			)
							_snake_case     	=	output.images

							_snake_case     	=	torch.Generator(device=__UpperCamelCase			).manual_seed(0			)
							_snake_case     	=	sd_pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=6.0,  num_inference_steps=2,  output_type="np",  return_dict=__UpperCamelCase,  )[0]

							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							_snake_case     	=	np.array([0.57_53, 0.61_14, 0.50_01, 0.50_34, 0.54_70, 0.47_29, 0.49_71, 0.48_67, 0.48_67]			)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1E-2






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	"cpu"  # ensure determinism for the device-dependent torch.Generator
							_snake_case     	=	self.dummy_cond_unet
							_snake_case     	=	PNDMScheduler(skip_prk_steps=__UpperCamelCase			)
							_snake_case     	=	self.dummy_vae
							_snake_case     	=	self.dummy_text_encoder
							_snake_case     	=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

							# make sure here that pndm scheduler skips prk
							_snake_case     	=	StableDiffusionPipeline(
							    unet=__UpperCamelCase,  scheduler=__UpperCamelCase,  vae=__UpperCamelCase,  text_encoder=__UpperCamelCase,  tokenizer=__UpperCamelCase,  safety_checker=__UpperCamelCase,  feature_extractor=self.dummy_extractor,  )
							_snake_case     	=	sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							_snake_case     	=	"A painting of a squirrel eating a burger"
							_snake_case     	=	torch.Generator(device=__UpperCamelCase			).manual_seed(0			)
							_snake_case     	=	sd_pipe([prompt],  generator=__UpperCamelCase,  guidance_scale=6.0,  num_inference_steps=2,  output_type="np"			)

							_snake_case     	=	output.images

							_snake_case     	=	torch.Generator(device=__UpperCamelCase			).manual_seed(0			)
							_snake_case     	=	sd_pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=6.0,  num_inference_steps=2,  output_type="np",  return_dict=__UpperCamelCase,  )[0]

							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							_snake_case     	=	np.array([0.51_22, 0.57_12, 0.48_25, 0.50_53, 0.56_46, 0.47_69, 0.51_79, 0.48_94, 0.49_94]			)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1E-2






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	StableDiffusionPipeline.from_pretrained(
							    "hf-internal-testing/tiny-stable-diffusion-lms-pipe",  safety_checker=__UpperCamelCase			)
							assert isinstance(__UpperCamelCase,  __UpperCamelCase			)
							assert isinstance(pipe.scheduler,  __UpperCamelCase			)
							assert pipe.safety_checker is None

							_snake_case     	=	pipe("example prompt",  num_inference_steps=2			).images[0]
							assert image is not None

							# check that there's no error when saving a pipeline with one of the models being None
							with tempfile.TemporaryDirectory() as tmpdirname:
									pipe.save_pretrained(__UpperCamelCase			)
									_snake_case     	=	StableDiffusionPipeline.from_pretrained(__UpperCamelCase			)

							# sanity check that the pipeline still works
							assert pipe.safety_checker is None
							_snake_case     	=	pipe("example prompt",  num_inference_steps=2			).images[0]
							assert image is not None






					@unittest.skipIf(torch_device != "cuda",  "This test requires a GPU"			)
					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	self.dummy_cond_unet
							_snake_case     	=	PNDMScheduler(skip_prk_steps=__UpperCamelCase			)
							_snake_case     	=	self.dummy_vae
							_snake_case     	=	self.dummy_text_encoder
							_snake_case     	=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)

							# put models in fp16
							_snake_case     	=	unet.half()
							_snake_case     	=	vae.half()
							_snake_case     	=	bert.half()

							# make sure here that pndm scheduler skips prk
							_snake_case     	=	StableDiffusionPipeline(
							    unet=__UpperCamelCase,  scheduler=__UpperCamelCase,  vae=__UpperCamelCase,  text_encoder=__UpperCamelCase,  tokenizer=__UpperCamelCase,  safety_checker=__UpperCamelCase,  feature_extractor=self.dummy_extractor,  )
							_snake_case     	=	sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							_snake_case     	=	"A painting of a squirrel eating a burger"
							_snake_case     	=	sd_pipe([prompt],  num_inference_steps=2,  output_type="np"			).images

							assert image.shape == (1, 64, 64, 3)





@nightly
@require_torch_gpu
class 					lowerCamelCase_							(	unittest.TestCase							):




					'''simple docstring'''






					def     __SCREAMING_SNAKE_CASE     (   self			):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	"cuda"
							_snake_case     	=	StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			)
							_snake_case     	=	pipe.to(__UpperCamelCase			)
							pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							_snake_case     	=	"a photo of a cat"
							_snake_case     	=	{
							    "editing_prompt": ["sunglasses"],
							    "reverse_editing_direction": [False],
							    "edit_warmup_steps": 10,
							    "edit_guidance_scale": 6,
							    "edit_threshold": 0.95,
							    "edit_momentum_scale": 0.5,
							    "edit_mom_beta": 0.6,
							}

							_snake_case     	=	3
							_snake_case     	=	7

							# no sega enabled
							_snake_case     	=	torch.Generator(__UpperCamelCase			)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.34_67_31_13,
							    0.38_49_27_33,
							    0.37_59_73_52,
							    0.34_08_63_35,
							    0.35_65_07_48,
							    0.35_57_92_05,
							    0.3_38_47_63,
							    0.34_34_02_36,
							    0.3_57_32_71,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2

							# with sega enabled
							# generator = torch.manual_seed(seed)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  **__UpperCamelCase,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.41_88_78_26,
							    0.37_72_87_66,
							    0.30_13_82_72,
							    0.41_41_63_35,
							    0.41_66_49_85,
							    0.36_28_33_92,
							    0.36_19_12_46,
							    0.43_36_44_65,
							    0.43_00_17_32,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	"cuda"
							_snake_case     	=	StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			)
							_snake_case     	=	pipe.to(__UpperCamelCase			)
							pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							_snake_case     	=	"an image of a crowded boulevard, realistic, 4k"
							_snake_case     	=	{
							    "editing_prompt": "crowd, crowded, people",
							    "reverse_editing_direction": True,
							    "edit_warmup_steps": 10,
							    "edit_guidance_scale": 8.3,
							    "edit_threshold": 0.9,
							    "edit_momentum_scale": 0.5,
							    "edit_mom_beta": 0.6,
							}

							_snake_case     	=	9
							_snake_case     	=	7

							# no sega enabled
							_snake_case     	=	torch.Generator(__UpperCamelCase			)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.43_49_79_98,
							    0.91_81_40_65,
							    0.7_54_07_39,
							    0.55_58_02_05,
							    0.8_46_72_65,
							    0.5_38_96_91,
							    0.62_57_45_06,
							    0.58_89_77_63,
							    0.50_92_67_57,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2

							# with sega enabled
							# generator = torch.manual_seed(seed)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  **__UpperCamelCase,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.3_08_97_19,
							    0.30_50_01_44,
							    0.29_01_60_42,
							    0.30_63_09_64,
							    0.32_56_87,
							    0.29_41_92_25,
							    0.2_90_80_91,
							    0.28_72_35_98,
							    0.27_69_62_94,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	"cuda"
							_snake_case     	=	StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"			)
							_snake_case     	=	pipe.to(__UpperCamelCase			)
							pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							_snake_case     	=	"a castle next to a river"
							_snake_case     	=	{
							    "editing_prompt": ["boat on a river, boat", "monet, impression, sunrise"],
							    "reverse_editing_direction": False,
							    "edit_warmup_steps": [15, 18],
							    "edit_guidance_scale": 6,
							    "edit_threshold": [0.9, 0.8],
							    "edit_momentum_scale": 0.5,
							    "edit_mom_beta": 0.6,
							}

							_snake_case     	=	48
							_snake_case     	=	7

							# no sega enabled
							_snake_case     	=	torch.Generator(__UpperCamelCase			)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.75_16_35_55,
							    0.76_03_71_45,
							    0.6_17_85,
							    0.9_18_96_73,
							    0.8_62_77_01,
							    0.85_18_96_94,
							    0.8_51_28_13,
							    0.87_01_20_76,
							    0.8_31_28_57,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2

							# with sega enabled
							# generator = torch.manual_seed(seed)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  **__UpperCamelCase,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.73_55_33_65,
							    0.7_53_72_71,
							    0.74_34_19_05,
							    0.66_48_03_56,
							    0.6_47_29_25,
							    0.63_03_94_16,
							    0.64_81_29_05,
							    0.6_74_97_17,
							    0.6_51_71_02,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2






					def     __SCREAMING_SNAKE_CASE     (   self			):
							_snake_case     	=	"cuda"
							_snake_case     	=	StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",  torch_dtype=torch.floataa			)
							_snake_case     	=	pipe.to(__UpperCamelCase			)
							pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							_snake_case     	=	"a photo of a cat"
							_snake_case     	=	{
							    "editing_prompt": ["sunglasses"],
							    "reverse_editing_direction": [False],
							    "edit_warmup_steps": 10,
							    "edit_guidance_scale": 6,
							    "edit_threshold": 0.95,
							    "edit_momentum_scale": 0.5,
							    "edit_mom_beta": 0.6,
							}

							_snake_case     	=	3
							_snake_case     	=	7

							# no sega enabled
							_snake_case     	=	torch.Generator(__UpperCamelCase			)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.34_88_76_95,
							    0.3_87_69_53,
							    0.3_75,
							    0.34_42_38_28,
							    0.3_58_15_43,
							    0.35_71_77_73,
							    0.3_38_37_89,
							    0.34_57_03_12,
							    0.35_93_75,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2

							# with sega enabled
							# generator = torch.manual_seed(seed)
							generator.manual_seed(__UpperCamelCase			)
							_snake_case     	=	pipe(
							    [prompt],  generator=__UpperCamelCase,  guidance_scale=__UpperCamelCase,  num_inference_steps=50,  output_type="np",  width=512,  height=512,  **__UpperCamelCase,  )

							_snake_case     	=	output.images
							_snake_case     	=	image[0, -3:, -3:, -1]
							_snake_case     	=	[
							    0.42_28_51_56,
							    0.36_91_40_62,
							    0.29_07_71_48,
							    0.42_04_10_16,
							    0.41_91_89_45,
							    0.35_49_80_47,
							    0.3_61_81_64,
							    0.4_42_38_28,
							    0.43_11_52_34,
							]

							assert image.shape == (1, 512, 512, 3)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2











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


_lowerCamelCase							:						Any							     =  logging.get_logger(__name__)

_lowerCamelCase							:						List[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'''






					lowercase__   =			"""bigbird_pegasus"""
					lowercase__   =			["""past_key_values"""]
					lowercase__   =			{
					    """num_attention_heads""": """encoder_attention_heads""",
					    """hidden_size""": """d_model""",
					    """attention_probs_dropout_prob""": """attention_dropout""",
					}






					def __init__(   self,  __UpperCamelCase=96_103,  __UpperCamelCase=4_096,  __UpperCamelCase=16,  __UpperCamelCase=4_096,  __UpperCamelCase=16,  __UpperCamelCase=16,  __UpperCamelCase=4_096,  __UpperCamelCase=16,  __UpperCamelCase=0.0,  __UpperCamelCase=0.0,  __UpperCamelCase=True,  __UpperCamelCase=True,  __UpperCamelCase="gelu_new",  __UpperCamelCase=1_024,  __UpperCamelCase=0.1,  __UpperCamelCase=0.0,  __UpperCamelCase=0.0,  __UpperCamelCase=0.02,  __UpperCamelCase=2,  __UpperCamelCase=0.0,  __UpperCamelCase=True,  __UpperCamelCase=0,  __UpperCamelCase=2,  __UpperCamelCase=1,  __UpperCamelCase="block_sparse",  __UpperCamelCase=64,  __UpperCamelCase=3,  __UpperCamelCase=False,  **__UpperCamelCase,  ):
							_snake_case     	=	vocab_size
							_snake_case     	=	max_position_embeddings
							_snake_case     	=	d_model
							_snake_case     	=	encoder_ffn_dim
							_snake_case     	=	encoder_layers
							_snake_case     	=	encoder_attention_heads
							_snake_case     	=	decoder_ffn_dim
							_snake_case     	=	decoder_layers
							_snake_case     	=	decoder_attention_heads
							_snake_case     	=	dropout
							_snake_case     	=	attention_dropout
							_snake_case     	=	activation_dropout
							_snake_case     	=	activation_function
							_snake_case     	=	init_std
							_snake_case     	=	encoder_layerdrop
							_snake_case     	=	decoder_layerdrop
							_snake_case     	=	classifier_dropout
							_snake_case     	=	use_cache
							_snake_case     	=	encoder_layers
							_snake_case     	=	scale_embedding  # scale factor will be sqrt(d_model) if True

							# extra config
							_snake_case     	=	attention_type
							_snake_case     	=	block_size
							_snake_case     	=	num_random_blocks
							_snake_case     	=	use_bias

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





class 					lowerCamelCase_							(	UpperCAmelCase_							):




					'''simple docstring'''






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

									if self.use_past:
											_snake_case     	=	{0: "batch"}
											_snake_case     	=	{0: "batch", 1: "past_decoder_sequence + sequence"}
									else:
											_snake_case     	=	{0: "batch", 1: "decoder_sequence"}
											_snake_case     	=	{0: "batch", 1: "decoder_sequence"}

									if self.use_past:
											self.fill_with_past_key_values_(__UpperCamelCase,  direction="inputs"			)
							elif self.task == "causal-lm":
									# TODO: figure this case out.
									_snake_case     	=	OrderedDict(
									    [
									        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
									        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
									    ]			)
									if self.use_past:
											_snake_case	,   _snake_case     	=	self.num_layers
											for i in range(__UpperCamelCase			):
													_snake_case     	=	{0: "batch", 2: "past_sequence + sequence"}
													_snake_case     	=	{0: "batch", 2: "past_sequence + sequence"}
							else:
									_snake_case     	=	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     __SCREAMING_SNAKE_CASE     (   self			):
							if self.task in ["default", "seq2seq-lm"]:
									_snake_case     	=	super().outputs
							else:
									_snake_case     	=	super(__UpperCamelCase,  self			).outputs
									if self.use_past:
											_snake_case	,   _snake_case     	=	self.num_layers
											for i in range(__UpperCamelCase			):
													_snake_case     	=	{0: "batch", 2: "past_sequence + sequence"}
													_snake_case     	=	{0: "batch", 2: "past_sequence + sequence"}
							return common_outputs






					def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase,  __UpperCamelCase = -1,  __UpperCamelCase = -1,  __UpperCamelCase = False,  __UpperCamelCase = None,  ):
							_snake_case     	=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)

							# Generate decoder inputs
							_snake_case     	=	seq_length if not self.use_past else 1
							_snake_case     	=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)
							_snake_case     	=	{f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
							_snake_case     	=	dict(**__UpperCamelCase,  **__UpperCamelCase			)

							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     	=	common_inputs["input_ids"].shape
									_snake_case     	=	common_inputs["decoder_input_ids"].shape[1]
									_snake_case	,   _snake_case     	=	self.num_attention_heads
									_snake_case     	=	(
									    batch,
									    num_encoder_attention_heads,
									    encoder_seq_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)
									_snake_case     	=	decoder_seq_length + 3
									_snake_case     	=	(
									    batch,
									    num_decoder_attention_heads,
									    decoder_past_length,
									    self._config.hidden_size // num_decoder_attention_heads,
									)

									_snake_case     	=	torch.cat(
									    [common_inputs["decoder_attention_mask"], torch.ones(__UpperCamelCase,  __UpperCamelCase			)],  dim=1			)

									_snake_case     	=	[]
									# If the number of encoder and decoder layers are present in the model configuration, both are considered
									_snake_case	,   _snake_case     	=	self.num_layers
									_snake_case     	=	min(__UpperCamelCase,  __UpperCamelCase			)
									_snake_case     	=	max(__UpperCamelCase,  __UpperCamelCase			) - min_num_layers
									_snake_case     	=	"encoder" if num_encoder_layers > num_decoder_layers else "decoder"

									for _ in range(__UpperCamelCase			):
											common_inputs["past_key_values"].append(
											    (
											        torch.zeros(__UpperCamelCase			),
											        torch.zeros(__UpperCamelCase			),
											        torch.zeros(__UpperCamelCase			),
											        torch.zeros(__UpperCamelCase			),
											    )			)
									# TODO: test this.
									_snake_case     	=	encoder_shape if remaining_side_name == "encoder" else decoder_shape
									for _ in range(__UpperCamelCase,  __UpperCamelCase			):
											common_inputs["past_key_values"].append((torch.zeros(__UpperCamelCase			), torch.zeros(__UpperCamelCase			))			)
							return common_inputs






					def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase,  __UpperCamelCase = -1,  __UpperCamelCase = -1,  __UpperCamelCase = False,  __UpperCamelCase = None,  ):
							_snake_case     	=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)

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

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






					def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase,  __UpperCamelCase = -1,  __UpperCamelCase = -1,  __UpperCamelCase = False,  __UpperCamelCase = 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
							_snake_case     	=	compute_effective_axis_dimension(
							    __UpperCamelCase,  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     	=	tokenizer.num_special_tokens_to_add(__UpperCamelCase			)
							_snake_case     	=	compute_effective_axis_dimension(
							    __UpperCamelCase,  fixed_dimension=OnnxConfig.default_fixed_sequence,  num_token_to_add=__UpperCamelCase			)

							# Generate dummy inputs according to compute batch and sequence
							_snake_case     	=	[" ".join([tokenizer.unk_token]			) * seq_length] * batch_size
							_snake_case     	=	dict(tokenizer(__UpperCamelCase,  return_tensors=__UpperCamelCase			)			)
							return common_inputs






					def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase,  __UpperCamelCase = -1,  __UpperCamelCase = -1,  __UpperCamelCase = False,  __UpperCamelCase = None,  ):
							if self.task in ["default", "seq2seq-lm"]:
									_snake_case     	=	self._generate_dummy_inputs_for_default_and_seqaseq_lm(
									    __UpperCamelCase,  batch_size=__UpperCamelCase,  seq_length=__UpperCamelCase,  is_pair=__UpperCamelCase,  framework=__UpperCamelCase			)

							elif self.task == "causal-lm":
									_snake_case     	=	self._generate_dummy_inputs_for_causal_lm(
									    __UpperCamelCase,  batch_size=__UpperCamelCase,  seq_length=__UpperCamelCase,  is_pair=__UpperCamelCase,  framework=__UpperCamelCase			)
							else:
									_snake_case     	=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    __UpperCamelCase,  batch_size=__UpperCamelCase,  seq_length=__UpperCamelCase,  is_pair=__UpperCamelCase,  framework=__UpperCamelCase			)

							return common_inputs






					def     __SCREAMING_SNAKE_CASE     (   self,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			):
							if self.task in ["default", "seq2seq-lm"]:
									_snake_case     	=	super()._flatten_past_key_values_(__UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)
							else:
									_snake_case     	=	super(__UpperCamelCase,  self			)._flatten_past_key_values_(
									    __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase,  __UpperCamelCase			)





