






import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ,     unittest.TestCase    ):


						'''simple docstring'''




						A   :      int						 =    UnCLIPPipeline
						A   :      str						 =    TEXT_TO_IMAGE_PARAMS - {
						    'negative_prompt',
						    'height',
						    'width',
						    'negative_prompt_embeds',
						    'guidance_scale',
						    'prompt_embeds',
						    'cross_attention_kwargs',
						}
						A   :      Dict						 =    TEXT_TO_IMAGE_BATCH_PARAMS
						A   :      Union[str, Any]						 =    [
						    'generator',
						    'return_dict',
						    'prior_num_inference_steps',
						    'decoder_num_inference_steps',
						    'super_res_num_inference_steps',
						]
						A   :      str						 =    False






						@property
						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								return 32






						@property
						def 					_lowerCAmelCase				(		self )       -> Any:
								return 32






						@property
						def 					_lowerCAmelCase				(		self )       -> Tuple:
								return self.time_input_dim






						@property
						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								return self.time_input_dim * 4






						@property
						def 					_lowerCAmelCase				(		self )       -> str:
								return 100






						@property
						def 					_lowerCAmelCase				(		self )       -> str:
								snake_case_ :							List[Any]         =				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )
								return tokenizer






						@property
						def 					_lowerCAmelCase				(		self )       -> int:
								torch.manual_seed(0 )
								snake_case_ :							Any         =				CLIPTextConfig(
								    bos_token_id=0						, eos_token_id=2						, hidden_size=self.text_embedder_hidden_size						, projection_dim=self.text_embedder_hidden_size						, intermediate_size=37						, layer_norm_eps=1e-05						, num_attention_heads=4						, num_hidden_layers=5						, pad_token_id=1						, vocab_size=1000						, )
								return CLIPTextModelWithProjection(_SCREAMING_SNAKE_CASE )






						@property
						def 					_lowerCAmelCase				(		self )       -> Tuple:
								torch.manual_seed(0 )

								snake_case_ :							Union[str, Any]         =				{
								    "num_attention_heads": 2,
								    "attention_head_dim": 12,
								    "embedding_dim": self.text_embedder_hidden_size,
								    "num_layers": 1,
								}

								snake_case_ :							Optional[int]         =				PriorTransformer(**_SCREAMING_SNAKE_CASE )
								return model






						@property
						def 					_lowerCAmelCase				(		self )       -> Any:
								torch.manual_seed(0 )

								snake_case_ :							List[Any]         =				{
								    "clip_embeddings_dim": self.text_embedder_hidden_size,
								    "time_embed_dim": self.time_embed_dim,
								    "cross_attention_dim": self.cross_attention_dim,
								}

								snake_case_ :							List[Any]         =				UnCLIPTextProjModel(**_SCREAMING_SNAKE_CASE )
								return model






						@property
						def 					_lowerCAmelCase				(		self )       -> Optional[int]:
								torch.manual_seed(0 )

								snake_case_ :							List[str]         =				{
								    "sample_size": 32,
								    # RGB in channels
								    "in_channels": 3,
								    # Out channels is double in channels because predicts mean and variance
								    "out_channels": 6,
								    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
								    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
								    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
								    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
								    "layers_per_block": 1,
								    "cross_attention_dim": self.cross_attention_dim,
								    "attention_head_dim": 4,
								    "resnet_time_scale_shift": "scale_shift",
								    "class_embed_type": "identity",
								}

								snake_case_ :							int         =				UNetaDConditionModel(**_SCREAMING_SNAKE_CASE )
								return model






						@property
						def 					_lowerCAmelCase				(		self )       -> Optional[int]:
								return {
								    "sample_size": 64,
								    "layers_per_block": 1,
								    "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
								    "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
								    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
								    "in_channels": 6,
								    "out_channels": 3,
								}






						@property
						def 					_lowerCAmelCase				(		self )       -> List[str]:
								torch.manual_seed(0 )

								snake_case_ :							int         =				UNetaDModel(**self.dummy_super_res_kwargs )
								return model






						@property
						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								# seeded differently to get different unet than `self.dummy_super_res_first`
								torch.manual_seed(1 )

								snake_case_ :							List[str]         =				UNetaDModel(**self.dummy_super_res_kwargs )
								return model






						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_ :							Dict         =				self.dummy_prior
								snake_case_ :							Any         =				self.dummy_decoder
								snake_case_ :							Optional[Any]         =				self.dummy_text_proj
								snake_case_ :							int         =				self.dummy_text_encoder
								snake_case_ :							Optional[int]         =				self.dummy_tokenizer
								snake_case_ :							List[Any]         =				self.dummy_super_res_first
								snake_case_ :							str         =				self.dummy_super_res_last

								snake_case_ :							Optional[Any]         =				UnCLIPScheduler(
								    variance_type="fixed_small_log"						, prediction_type="sample"						, num_train_timesteps=1000						, clip_sample_range=5.0						, )

								snake_case_ :							Tuple         =				UnCLIPScheduler(
								    variance_type="learned_range"						, prediction_type="epsilon"						, num_train_timesteps=1000						, )

								snake_case_ :							Tuple         =				UnCLIPScheduler(
								    variance_type="fixed_small_log"						, prediction_type="epsilon"						, num_train_timesteps=1000						, )

								snake_case_ :							Any         =				{
								    "prior": prior,
								    "decoder": decoder,
								    "text_proj": text_proj,
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "super_res_first": super_res_first,
								    "super_res_last": super_res_last,
								    "prior_scheduler": prior_scheduler,
								    "decoder_scheduler": decoder_scheduler,
								    "super_res_scheduler": super_res_scheduler,
								}

								return components






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE=0 )       -> Union[str, Any]:
								if str(_SCREAMING_SNAKE_CASE ).startswith("mps" ):
										snake_case_ :							Dict         =				torch.manual_seed(_SCREAMING_SNAKE_CASE )
								else:
										snake_case_ :							Union[str, Any]         =				torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(_SCREAMING_SNAKE_CASE )
								snake_case_ :							List[str]         =				{
								    "prompt": "horse",
								    "generator": generator,
								    "prior_num_inference_steps": 2,
								    "decoder_num_inference_steps": 2,
								    "super_res_num_inference_steps": 2,
								    "output_type": "numpy",
								}
								return inputs






						def 					_lowerCAmelCase				(		self )       -> int:
								snake_case_ :							Optional[int]         =				"cpu"

								snake_case_ :							List[str]         =				self.get_dummy_components()

								snake_case_ :							int         =				self.pipeline_class(**_SCREAMING_SNAKE_CASE )
								snake_case_ :							Tuple         =				pipe.to(_SCREAMING_SNAKE_CASE )

								pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

								snake_case_ :							Any         =				pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE ) )
								snake_case_ :							str         =				output.images

								snake_case_ :							int         =				pipe(
								    **self.get_dummy_inputs(_SCREAMING_SNAKE_CASE )						, return_dict=_SCREAMING_SNAKE_CASE						, )[0]

								snake_case_ :							int         =				image[0, -3:, -3:, -1]
								snake_case_ :							Optional[Any]         =				image_from_tuple[0, -3:, -3:, -1]

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

								snake_case_ :							Union[str, Any]         =				np.array(
								    [
								        0.9997,
								        0.9988,
								        0.0028,
								        0.9997,
								        0.9984,
								        0.9965,
								        0.0029,
								        0.9986,
								        0.0025,
								    ] )

								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 					_lowerCAmelCase				(		self )       -> Optional[int]:
								snake_case_ :							Union[str, Any]         =				torch.device("cpu" )




								class 						UpperCAmelCase_  :


														'''simple docstring'''




														A   :      Union[str, Any]						 =    1




								snake_case_ :							Any         =				self.get_dummy_components()

								snake_case_ :							Optional[Any]         =				self.pipeline_class(**_SCREAMING_SNAKE_CASE )
								snake_case_ :							str         =				pipe.to(_SCREAMING_SNAKE_CASE )

								snake_case_ :							Any         =				components["prior"]
								snake_case_ :							str         =				components["decoder"]
								snake_case_ :							List[str]         =				components["super_res_first"]
								snake_case_ :							str         =				components["tokenizer"]
								snake_case_ :							Any         =				components["text_encoder"]

								snake_case_ :							List[str]         =				torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(0 )
								snake_case_ :							Dict         =				prior.dtype
								snake_case_ :							Tuple         =				1

								snake_case_ :							Optional[int]         =				(batch_size, prior.config.embedding_dim)
								snake_case_ :							List[Any]         =				pipe.prepare_latents(
								    _SCREAMING_SNAKE_CASE						, dtype=_SCREAMING_SNAKE_CASE						, device=_SCREAMING_SNAKE_CASE						, generator=_SCREAMING_SNAKE_CASE						, latents=_SCREAMING_SNAKE_CASE						, scheduler=DummyScheduler() )
								snake_case_ :							str         =				(batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
								snake_case_ :							List[Any]         =				pipe.prepare_latents(
								    _SCREAMING_SNAKE_CASE						, dtype=_SCREAMING_SNAKE_CASE						, device=_SCREAMING_SNAKE_CASE						, generator=_SCREAMING_SNAKE_CASE						, latents=_SCREAMING_SNAKE_CASE						, scheduler=DummyScheduler() )

								snake_case_ :							Tuple         =				(
								    batch_size,
								    super_res_first.config.in_channels // 2,
								    super_res_first.config.sample_size,
								    super_res_first.config.sample_size,
								)
								snake_case_ :							int         =				pipe.prepare_latents(
								    _SCREAMING_SNAKE_CASE						, dtype=_SCREAMING_SNAKE_CASE						, device=_SCREAMING_SNAKE_CASE						, generator=_SCREAMING_SNAKE_CASE						, latents=_SCREAMING_SNAKE_CASE						, scheduler=DummyScheduler() )

								pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

								snake_case_ :							int         =				"this is a prompt example"

								snake_case_ :							Optional[int]         =				torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(0 )
								snake_case_ :							str         =				pipe(
								    [prompt]						, generator=_SCREAMING_SNAKE_CASE						, prior_num_inference_steps=2						, decoder_num_inference_steps=2						, super_res_num_inference_steps=2						, prior_latents=_SCREAMING_SNAKE_CASE						, decoder_latents=_SCREAMING_SNAKE_CASE						, super_res_latents=_SCREAMING_SNAKE_CASE						, output_type="np"						, )
								snake_case_ :							List[Any]         =				output.images

								snake_case_ :							Tuple         =				tokenizer(
								    _SCREAMING_SNAKE_CASE						, padding="max_length"						, max_length=tokenizer.model_max_length						, return_tensors="pt"						, )
								snake_case_ :							List[Any]         =				text_encoder(text_inputs.input_ids )
								snake_case_ :							str         =				text_inputs.attention_mask

								snake_case_ :							List[Any]         =				torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(0 )
								snake_case_ :							Optional[int]         =				pipe(
								    generator=_SCREAMING_SNAKE_CASE						, prior_num_inference_steps=2						, decoder_num_inference_steps=2						, super_res_num_inference_steps=2						, prior_latents=_SCREAMING_SNAKE_CASE						, decoder_latents=_SCREAMING_SNAKE_CASE						, super_res_latents=_SCREAMING_SNAKE_CASE						, text_model_output=_SCREAMING_SNAKE_CASE						, text_attention_mask=_SCREAMING_SNAKE_CASE						, output_type="np"						, )[0]

								# make sure passing text embeddings manually is identical
								assert np.abs(image - image_from_text ).max() < 1e-4






						@skip_mps
						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								snake_case_ :							List[str]         =				torch_device == "cpu"

								self._test_attention_slicing_forward_pass(test_max_difference=_SCREAMING_SNAKE_CASE						, expected_max_diff=0.01 )






						@skip_mps
						def 					_lowerCAmelCase				(		self )       -> Tuple:
								snake_case_ :							str         =				torch_device == "cpu"
								snake_case_ :							str         =				True
								snake_case_ :							Union[str, Any]         =				[
								    "prior_num_inference_steps",
								    "decoder_num_inference_steps",
								    "super_res_num_inference_steps",
								]

								self._test_inference_batch_single_identical(
								    test_max_difference=_SCREAMING_SNAKE_CASE						, relax_max_difference=_SCREAMING_SNAKE_CASE						, additional_params_copy_to_batched_inputs=_SCREAMING_SNAKE_CASE						, )






						def 					_lowerCAmelCase				(		self )       -> Dict:
								snake_case_ :							List[str]         =				[
								    "prior_num_inference_steps",
								    "decoder_num_inference_steps",
								    "super_res_num_inference_steps",
								]

								if torch_device == "mps":
										# TODO: MPS errors with larger batch sizes
										snake_case_ :							List[str]         =				[2, 3]
										self._test_inference_batch_consistent(
										    batch_sizes=_SCREAMING_SNAKE_CASE						, additional_params_copy_to_batched_inputs=_SCREAMING_SNAKE_CASE						, )
								else:
										self._test_inference_batch_consistent(
										    additional_params_copy_to_batched_inputs=_SCREAMING_SNAKE_CASE )






						@skip_mps
						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								return super().test_dict_tuple_outputs_equivalent()






						@skip_mps
						def 					_lowerCAmelCase				(		self )       -> List[Any]:
								return super().test_save_load_local()






						@skip_mps
						def 					_lowerCAmelCase				(		self )       -> Any:
								return super().test_save_load_optional_components()




@nightly
class 						UpperCAmelCase_  (		unittest.TestCase    ):


						'''simple docstring'''




						def 					_lowerCAmelCase				(		self )       -> Optional[Any]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






						def 					_lowerCAmelCase				(		self )       -> str:
								snake_case_ :							Union[str, Any]         =				load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/unclip/karlo_v1_alpha_horse_cpu.npy" )

								snake_case_ :							int         =				UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha" )
								pipeline.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

								snake_case_ :							Union[str, Any]         =				torch.manual_seed(0 )
								snake_case_ :							str         =				pipeline(
								    "horse"						, num_images_per_prompt=1						, generator=_SCREAMING_SNAKE_CASE						, output_type="np"						, )

								snake_case_ :							Any         =				output.images[0]

								assert image.shape == (256, 256, 3)
								assert np.abs(expected_image - image ).max() < 1e-1




@slow
@require_torch_gpu
class 						UpperCAmelCase_  (		unittest.TestCase    ):


						'''simple docstring'''




						def 					_lowerCAmelCase				(		self )       -> Optional[int]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






						def 					_lowerCAmelCase				(		self )       -> List[str]:
								snake_case_ :							Optional[int]         =				load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/unclip/karlo_v1_alpha_horse_fp16.npy" )

								snake_case_ :							str         =				UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"						, torch_dtype=torch.floataa )
								snake_case_ :							str         =				pipeline.to(_SCREAMING_SNAKE_CASE )
								pipeline.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

								snake_case_ :							Dict         =				torch.Generator(device="cpu" ).manual_seed(0 )
								snake_case_ :							Dict         =				pipeline(
								    "horse"						, generator=_SCREAMING_SNAKE_CASE						, output_type="np"						, )

								snake_case_ :							int         =				output.images[0]

								assert image.shape == (256, 256, 3)

								assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )






						def 					_lowerCAmelCase				(		self )       -> Union[str, Any]:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								snake_case_ :							Dict         =				UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"						, torch_dtype=torch.floataa )
								snake_case_ :							List[str]         =				pipe.to(_SCREAMING_SNAKE_CASE )
								pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )
								pipe.enable_attention_slicing()
								pipe.enable_sequential_cpu_offload()

								snake_case_ :							List[str]         =				pipe(
								    "horse"						, num_images_per_prompt=1						, prior_num_inference_steps=2						, decoder_num_inference_steps=2						, super_res_num_inference_steps=2						, output_type="np"						, )

								snake_case_ :							Dict         =				torch.cuda.max_memory_allocated()
								# make sure that less than 7 GB is allocated
								assert mem_bytes < 7 * 10**9








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


lowercase     :							List[str]	     =			logging.get_logger(__name__)

lowercase     :							int	     =			{
    '''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 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ):


						'''simple docstring'''




						A   :      List[str]						 =    'bigbird_pegasus'
						A   :      Optional[int]						 =    ['past_key_values']
						A   :      Tuple						 =    {
						    'num_attention_heads': 'encoder_attention_heads',
						    'hidden_size': 'd_model',
						    'attention_probs_dropout_prob': 'attention_dropout',
						}






						def __init__(		self						, _SCREAMING_SNAKE_CASE=9_6103						, _SCREAMING_SNAKE_CASE=4096						, _SCREAMING_SNAKE_CASE=16						, _SCREAMING_SNAKE_CASE=4096						, _SCREAMING_SNAKE_CASE=16						, _SCREAMING_SNAKE_CASE=16						, _SCREAMING_SNAKE_CASE=4096						, _SCREAMING_SNAKE_CASE=16						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=True						, _SCREAMING_SNAKE_CASE=True						, _SCREAMING_SNAKE_CASE="gelu_new"						, _SCREAMING_SNAKE_CASE=1024						, _SCREAMING_SNAKE_CASE=0.1						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=0.02						, _SCREAMING_SNAKE_CASE=2						, _SCREAMING_SNAKE_CASE=0.0						, _SCREAMING_SNAKE_CASE=True						, _SCREAMING_SNAKE_CASE=0						, _SCREAMING_SNAKE_CASE=2						, _SCREAMING_SNAKE_CASE=1						, _SCREAMING_SNAKE_CASE="block_sparse"						, _SCREAMING_SNAKE_CASE=64						, _SCREAMING_SNAKE_CASE=3						, _SCREAMING_SNAKE_CASE=False						, **_SCREAMING_SNAKE_CASE						, )       -> Dict:
								snake_case_ :							str         =				vocab_size
								snake_case_ :							Tuple         =				max_position_embeddings
								snake_case_ :							List[Any]         =				d_model
								snake_case_ :							List[str]         =				encoder_ffn_dim
								snake_case_ :							List[str]         =				encoder_layers
								snake_case_ :							int         =				encoder_attention_heads
								snake_case_ :							Tuple         =				decoder_ffn_dim
								snake_case_ :							List[Any]         =				decoder_layers
								snake_case_ :							Optional[int]         =				decoder_attention_heads
								snake_case_ :							Any         =				dropout
								snake_case_ :							Tuple         =				attention_dropout
								snake_case_ :							List[str]         =				activation_dropout
								snake_case_ :							Any         =				activation_function
								snake_case_ :							Any         =				init_std
								snake_case_ :							Optional[int]         =				encoder_layerdrop
								snake_case_ :							Optional[Any]         =				decoder_layerdrop
								snake_case_ :							Tuple         =				classifier_dropout
								snake_case_ :							str         =				use_cache
								snake_case_ :							Dict         =				encoder_layers
								snake_case_ :							List[str]         =				scale_embedding  # scale factor will be sqrt(d_model) if True

								# extra config
								snake_case_ :							Union[str, Any]         =				attention_type
								snake_case_ :							Tuple         =				block_size
								snake_case_ :							Any         =				num_random_blocks
								snake_case_ :							List[str]         =				use_bias

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




class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ):


						'''simple docstring'''




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

										if self.use_past:
												snake_case_ :							List[str]         =				{0: "batch"}
												snake_case_ :							Dict         =				{0: "batch", 1: "past_decoder_sequence + sequence"}
										else:
												snake_case_ :							Any         =				{0: "batch", 1: "decoder_sequence"}
												snake_case_ :							int         =				{0: "batch", 1: "decoder_sequence"}

										if self.use_past:
												self.fill_with_past_key_values_(_SCREAMING_SNAKE_CASE						, direction="inputs" )
								elif self.task == "causal-lm":
										# TODO: figure this case out.
										snake_case_ :							Union[str, Any]         =				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(_SCREAMING_SNAKE_CASE ):
														snake_case_ :							int         =				{0: "batch", 2: "past_sequence + sequence"}
														snake_case_ :							Dict         =				{0: "batch", 2: "past_sequence + sequence"}
								else:
										snake_case_ :							Union[str, 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 					_lowerCAmelCase				(		self )       -> Mapping[str, Mapping[int, str]]:
								if self.task in ["default", "seq2seq-lm"]:
										snake_case_ :							str         =				super().outputs
								else:
										snake_case_ :							Optional[Any]         =				super(_SCREAMING_SNAKE_CASE						, self ).outputs
										if self.use_past:
												snake_case_       ,				snake_case_ :							Optional[int]         =				self.num_layers
												for i in range(_SCREAMING_SNAKE_CASE ):
														snake_case_ :							Tuple         =				{0: "batch", 2: "past_sequence + sequence"}
														snake_case_ :							Union[str, Any]         =				{0: "batch", 2: "past_sequence + sequence"}
								return common_outputs






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = None						, )       -> Mapping[str, Any]:
								snake_case_ :							List[str]         =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
								    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )

								# Generate decoder inputs
								snake_case_ :							Optional[int]         =				seq_length if not self.use_past else 1
								snake_case_ :							Dict         =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
								    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
								snake_case_ :							List[Any]         =				{f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
								snake_case_ :							str         =				dict(**_SCREAMING_SNAKE_CASE						, **_SCREAMING_SNAKE_CASE )

								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_ :							str         =				common_inputs["input_ids"].shape
										snake_case_ :							List[str]         =				common_inputs["decoder_input_ids"].shape[1]
										snake_case_       ,				snake_case_ :							Tuple         =				self.num_attention_heads
										snake_case_ :							List[Any]         =				(
										    batch,
										    num_encoder_attention_heads,
										    encoder_seq_length,
										    self._config.hidden_size // num_encoder_attention_heads,
										)
										snake_case_ :							Optional[Any]         =				decoder_seq_length + 3
										snake_case_ :							Any         =				(
										    batch,
										    num_decoder_attention_heads,
										    decoder_past_length,
										    self._config.hidden_size // num_decoder_attention_heads,
										)

										snake_case_ :							Optional[Any]         =				torch.cat(
										    [common_inputs["decoder_attention_mask"], torch.ones(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )]						, dim=1 )

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

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = None						, )       -> Mapping[str, Any]:
								snake_case_ :							Union[str, Any]         =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
								    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )

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

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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = None						, )       -> Mapping[str, Any]:
								# 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_ :							List[str]         =				compute_effective_axis_dimension(
								    _SCREAMING_SNAKE_CASE						, 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_ :							Any         =				tokenizer.num_special_tokens_to_add(_SCREAMING_SNAKE_CASE )
								snake_case_ :							int         =				compute_effective_axis_dimension(
								    _SCREAMING_SNAKE_CASE						, fixed_dimension=OnnxConfig.default_fixed_sequence						, num_token_to_add=_SCREAMING_SNAKE_CASE )

								# Generate dummy inputs according to compute batch and sequence
								snake_case_ :							List[Any]         =				[" ".join([tokenizer.unk_token] ) * seq_length] * batch_size
								snake_case_ :							Tuple         =				dict(tokenizer(_SCREAMING_SNAKE_CASE						, return_tensors=_SCREAMING_SNAKE_CASE ) )
								return common_inputs






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = -1						, _SCREAMING_SNAKE_CASE = False						, _SCREAMING_SNAKE_CASE = None						, )       -> Mapping[str, Any]:
								if self.task in ["default", "seq2seq-lm"]:
										snake_case_ :							List[Any]         =				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
										    _SCREAMING_SNAKE_CASE						, batch_size=_SCREAMING_SNAKE_CASE						, seq_length=_SCREAMING_SNAKE_CASE						, is_pair=_SCREAMING_SNAKE_CASE						, framework=_SCREAMING_SNAKE_CASE )

								elif self.task == "causal-lm":
										snake_case_ :							Optional[Any]         =				self._generate_dummy_inputs_for_causal_lm(
										    _SCREAMING_SNAKE_CASE						, batch_size=_SCREAMING_SNAKE_CASE						, seq_length=_SCREAMING_SNAKE_CASE						, is_pair=_SCREAMING_SNAKE_CASE						, framework=_SCREAMING_SNAKE_CASE )
								else:
										snake_case_ :							Any         =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    _SCREAMING_SNAKE_CASE						, batch_size=_SCREAMING_SNAKE_CASE						, seq_length=_SCREAMING_SNAKE_CASE						, is_pair=_SCREAMING_SNAKE_CASE						, framework=_SCREAMING_SNAKE_CASE )

								return common_inputs






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Any:
								if self.task in ["default", "seq2seq-lm"]:
										snake_case_ :							int         =				super()._flatten_past_key_values_(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
								else:
										snake_case_ :							Optional[Any]         =				super(_SCREAMING_SNAKE_CASE						, self )._flatten_past_key_values_(
										    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
