





import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device







def  lowerCAmelCase__(		lowercase    :      nn.Module			)	->      Optional[Any]:
  __snake_case				:			Optional[Any]            =    {}
  for name in unet.attn_processors.keys():
    __snake_case				:			str            =    None if name.endswith("attn1.processor"			) else unet.config.cross_attention_dim
    if name.startswith("mid_block"			):
      __snake_case				:			Optional[Any]            =    unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"			):
      __snake_case				:			int            =    int(name[len("up_blocks."			)]			)
      __snake_case				:			List[Any]            =    list(reversed(unet.config.block_out_channels			)			)[block_id]
    elif name.startswith("down_blocks"			):
      __snake_case				:			List[str]            =    int(name[len("down_blocks."			)]			)
      __snake_case				:			List[Any]            =    unet.config.block_out_channels[block_id]
    __snake_case				:			Dict            =    (
        LoRAAttnProcessor2_0 if hasattr(lowercase				,       "scaled_dot_product_attention"			) else LoRAAttnProcessor
    )
    __snake_case				:			int            =    lora_attn_processor_class(
        hidden_size=lowercase				,       cross_attention_dim=lowercase			)
  __snake_case				:			str            =    AttnProcsLayers(lowercase			)
  return lora_attn_procs, unet_lora_layers







def  lowerCAmelCase__(		lowercase    :      nn.Module			)	->      Union[str, Any]:
  __snake_case				:			Union[str, Any]            =    {}
  __snake_case				:			List[Any]            =    (
      LoRAAttnProcessor2_0 if hasattr(lowercase				,       "scaled_dot_product_attention"			) else LoRAAttnProcessor
  )
  for name, module in text_encoder.named_modules():
    if name.endswith(lowercase			):
      __snake_case				:			str            =    lora_attn_processor_class(
          hidden_size=module.out_proj.out_features				,       cross_attention_dim=lowercase			)
  return text_lora_attn_procs







def  lowerCAmelCase__(		lowercase    :      nn.Module			)	->      List[Any]:
  __snake_case				:			Dict            =    create_text_encoder_lora_attn_procs(lowercase			)
  __snake_case				:			Tuple            =    AttnProcsLayers(lowercase			)
  return text_encoder_lora_layers







def  lowerCAmelCase__(		lowercase    :      List[Any]				,       lowercase    :      int=False			)	->      Union[str, Any]:
  for _, attn_proc in text_lora_attn_procs.items():
    # set up.weights
    for layer_name, layer_module in attn_proc.named_modules():
      if layer_name.endswith("_lora"			):
        __snake_case				:			int            =    (
            torch.randn_like(layer_module.up.weight			)
            if randn_weight
            else torch.zeros_like(layer_module.up.weight			)
        )
        __snake_case				:			List[Any]            =    torch.nn.Parameter(lowercase			)




class 			_lowerCamelCase	(      unittest.TestCase      ):



      """simple docstring"""






      def   UpperCAmelCase						(	self				)    ->				Union[str, Any]:






        '''simple docstring'''






        torch.manual_seed(0				)
        __snake_case				:			Optional[int]            =    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			,  )
        __snake_case				:			str            =    DDIMScheduler(
            beta_start=0.00_085			,  beta_end=0.012			,  beta_schedule="scaled_linear"			,  clip_sample=UpperCAmelCase			,  set_alpha_to_one=UpperCAmelCase			,  steps_offset=1			,  )
        torch.manual_seed(0				)
        __snake_case				:			Dict            =    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			,  )
        __snake_case				:			Optional[Any]            =    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=1000			,  )
        __snake_case				:			Optional[int]            =    CLIPTextModel(UpperCAmelCase				)
        __snake_case				:			Any            =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)

        __snake_case				:			int            =    create_unet_lora_layers(UpperCAmelCase				)
        __snake_case				:			int            =    create_text_encoder_lora_layers(UpperCAmelCase				)

        __snake_case				:			Dict            =    {
            "unet": unet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        __snake_case				:			Optional[int]            =    {
            "unet_lora_layers": unet_lora_layers,
            "text_encoder_lora_layers": text_encoder_lora_layers,
            "unet_lora_attn_procs": unet_lora_attn_procs,
        }
        return pipeline_components, lora_components







      def   UpperCAmelCase						(	self				)    ->				Tuple:






        '''simple docstring'''






        __snake_case				:			List[Any]            =    1
        __snake_case				:			Optional[int]            =    10
        __snake_case				:			Dict            =    4
        __snake_case				:			Any            =    (32, 32)

        __snake_case				:			str            =    torch.manual_seed(0				)
        __snake_case				:			Optional[int]            =    floats_tensor((batch_size, num_channels) + sizes				)
        __snake_case				:			Dict            =    torch.randint(1			,  UpperCAmelCase			,  size=(batch_size, sequence_length)			,  generator=UpperCAmelCase				)

        __snake_case				:			Dict            =    {
            "prompt": "A painting of a squirrel eating a burger",
            "generator": generator,
            "num_inference_steps": 2,
            "guidance_scale": 6.0,
            "output_type": "numpy",
        }

        return noise, input_ids, pipeline_inputs

        # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb







      def   UpperCAmelCase						(	self				)    ->				int:






        '''simple docstring'''






        __snake_case				:			Union[str, Any]            =    77

        __snake_case				:			Union[str, Any]            =    torch.randint(2			,  56			,  size=(1, max_seq_length)			,  generator=torch.manual_seed(0				)				)

        __snake_case				:			Optional[int]            =    {}
        __snake_case				:			str            =    inputs
        return prepared_inputs







      def   UpperCAmelCase						(	self			,  UpperCAmelCase				)    ->				Tuple:






        '''simple docstring'''






        __snake_case				:			Dict            =    self.get_dummy_components()
        LoraLoaderMixin.save_lora_weights(
            save_directory=UpperCAmelCase			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
        self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase			,  "pytorch_lora_weights.bin"				)				)				)







      def   UpperCAmelCase						(	self				)    ->				int:






        '''simple docstring'''






        __snake_case				:			Tuple            =    self.get_dummy_components()
        __snake_case				:			List[str]            =    StableDiffusionPipeline(**UpperCAmelCase				)
        __snake_case				:			Any            =    sd_pipe.to(UpperCAmelCase				)
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			List[str]            =    self.get_dummy_inputs()

        __snake_case				:			str            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			int            =    original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
          LoraLoaderMixin.save_lora_weights(
              save_directory=UpperCAmelCase			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
          self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase			,  "pytorch_lora_weights.bin"				)				)				)
          sd_pipe.load_lora_weights(UpperCAmelCase				)

        __snake_case				:			int            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			int            =    lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(UpperCAmelCase				)			,  torch.from_numpy(UpperCAmelCase				)				)				)







      def   UpperCAmelCase						(	self				)    ->				str:






        '''simple docstring'''






        __snake_case				:			Tuple            =    self.get_dummy_components()
        __snake_case				:			List[Any]            =    StableDiffusionPipeline(**UpperCAmelCase				)
        __snake_case				:			Optional[int]            =    sd_pipe.to(UpperCAmelCase				)
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			List[Any]            =    self.get_dummy_inputs()

        __snake_case				:			int            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			int            =    original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
          LoraLoaderMixin.save_lora_weights(
              save_directory=UpperCAmelCase			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  safe_serialization=UpperCAmelCase			,  )
          self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase			,  "pytorch_lora_weights.safetensors"				)				)				)
          sd_pipe.load_lora_weights(UpperCAmelCase				)

        __snake_case				:			Any            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			Union[str, Any]            =    lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(UpperCAmelCase				)			,  torch.from_numpy(UpperCAmelCase				)				)				)







      def   UpperCAmelCase						(	self				)    ->				Tuple:






        '''simple docstring'''






        __snake_case				:			List[str]            =    self.get_dummy_components()
        __snake_case				:			Any            =    lora_components["unet_lora_attn_procs"]
        __snake_case				:			int            =    StableDiffusionPipeline(**UpperCAmelCase				)
        __snake_case				:			str            =    sd_pipe.to(UpperCAmelCase				)
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			Any            =    self.get_dummy_inputs()

        __snake_case				:			Dict            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			int            =    original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
          __snake_case				:			Union[str, Any]            =    sd_pipe.unet
          unet.set_attn_processor(UpperCAmelCase				)
          unet.save_attn_procs(UpperCAmelCase				)
          self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase			,  "pytorch_lora_weights.bin"				)				)				)
          sd_pipe.load_lora_weights(UpperCAmelCase				)

        __snake_case				:			Tuple            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			Tuple            =    lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(UpperCAmelCase				)			,  torch.from_numpy(UpperCAmelCase				)				)				)







      def   UpperCAmelCase						(	self				)    ->				int:






        '''simple docstring'''






        __snake_case				:			int            =    self.get_dummy_components()
        __snake_case				:			List[str]            =    StableDiffusionPipeline(**UpperCAmelCase				)

        __snake_case				:			Optional[Any]            =    self.get_dummy_tokens()

        # inference without lora
        __snake_case				:			Any            =    pipe.text_encoder(**UpperCAmelCase				)[0]
        assert outputs_without_lora.shape == (1, 77, 32)

        # create lora_attn_procs with zeroed out up.weights
        __snake_case				:			Union[str, Any]            =    create_text_encoder_lora_attn_procs(pipe.text_encoder				)
        set_lora_up_weights(UpperCAmelCase			,  randn_weight=UpperCAmelCase				)

        # monkey patch
        pipe._modify_text_encoder(UpperCAmelCase				)

        # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
        del text_attn_procs
        gc.collect()

        # inference with lora
        __snake_case				:			Any            =    pipe.text_encoder(**UpperCAmelCase				)[0]
        assert outputs_with_lora.shape == (1, 77, 32)

        assert torch.allclose(
            UpperCAmelCase			,  UpperCAmelCase				), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

        # create lora_attn_procs with randn up.weights
        __snake_case				:			Union[str, Any]            =    create_text_encoder_lora_attn_procs(pipe.text_encoder				)
        set_lora_up_weights(UpperCAmelCase			,  randn_weight=UpperCAmelCase				)

        # monkey patch
        pipe._modify_text_encoder(UpperCAmelCase				)

        # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
        del text_attn_procs
        gc.collect()

        # inference with lora
        __snake_case				:			Optional[int]            =    pipe.text_encoder(**UpperCAmelCase				)[0]
        assert outputs_with_lora.shape == (1, 77, 32)

        assert not torch.allclose(
            UpperCAmelCase			,  UpperCAmelCase				), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"







      def   UpperCAmelCase						(	self				)    ->				Dict:






        '''simple docstring'''






        __snake_case				:			int            =    self.get_dummy_components()
        __snake_case				:			Tuple            =    StableDiffusionPipeline(**UpperCAmelCase				)

        __snake_case				:			List[Any]            =    self.get_dummy_tokens()

        # inference without lora
        __snake_case				:			str            =    pipe.text_encoder(**UpperCAmelCase				)[0]
        assert outputs_without_lora.shape == (1, 77, 32)

        # create lora_attn_procs with randn up.weights
        __snake_case				:			str            =    create_text_encoder_lora_attn_procs(pipe.text_encoder				)
        set_lora_up_weights(UpperCAmelCase			,  randn_weight=UpperCAmelCase				)

        # monkey patch
        pipe._modify_text_encoder(UpperCAmelCase				)

        # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
        del text_attn_procs
        gc.collect()

        # inference with lora
        __snake_case				:			Optional[Any]            =    pipe.text_encoder(**UpperCAmelCase				)[0]
        assert outputs_with_lora.shape == (1, 77, 32)

        assert not torch.allclose(
            UpperCAmelCase			,  UpperCAmelCase				), "lora outputs should be different to without lora outputs"

        # remove monkey patch
        pipe._remove_text_encoder_monkey_patch()

        # inference with removed lora
        __snake_case				:			Optional[Any]            =    pipe.text_encoder(**UpperCAmelCase				)[0]
        assert outputs_without_lora_removed.shape == (1, 77, 32)

        assert torch.allclose(
            UpperCAmelCase			,  UpperCAmelCase				), "remove lora monkey patch should restore the original outputs"







      def   UpperCAmelCase						(	self				)    ->				int:






        '''simple docstring'''






        __snake_case				:			Dict            =    self.get_dummy_components()
        __snake_case				:			List[Any]            =    StableDiffusionPipeline(**UpperCAmelCase				)
        __snake_case				:			Union[str, Any]            =    sd_pipe.to(UpperCAmelCase				)
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			List[str]            =    self.get_dummy_inputs()

        with tempfile.TemporaryDirectory() as tmpdirname:
          LoraLoaderMixin.save_lora_weights(
              save_directory=UpperCAmelCase			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
          self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase			,  "pytorch_lora_weights.bin"				)				)				)
          sd_pipe.load_lora_weights(UpperCAmelCase				)

        __snake_case				:			Optional[int]            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			Dict            =    lora_images[0, -3:, -3:, -1]

        __snake_case				:			Dict            =    sd_pipe(**UpperCAmelCase			,  cross_attention_kwargs={"scale": 0.5}				).images
        __snake_case				:			Dict            =    lora_images_with_scale[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(
            torch.allclose(torch.from_numpy(UpperCAmelCase				)			,  torch.from_numpy(UpperCAmelCase				)				)				)







      def   UpperCAmelCase						(	self				)    ->				List[Any]:






        '''simple docstring'''






        with tempfile.TemporaryDirectory() as tmpdirname:
          self.create_lora_weight_file(UpperCAmelCase				)

          __snake_case				:			Any            =    self.get_dummy_components()
          __snake_case				:			List[str]            =    StableDiffusionPipeline(**UpperCAmelCase				)
          __snake_case				:			Dict            =    sd_pipe.to(UpperCAmelCase				)
          sd_pipe.set_progress_bar_config(disable=UpperCAmelCase				)

          # check if vanilla attention processors are used
          for _, module in sd_pipe.unet.named_modules():
            if isinstance(UpperCAmelCase			,  UpperCAmelCase				):
              self.assertIsInstance(module.processor			,  (AttnProcessor, AttnProcessor2_0)				)

            # load LoRA weight file
          sd_pipe.load_lora_weights(UpperCAmelCase				)

          # check if lora attention processors are used
          for _, module in sd_pipe.unet.named_modules():
            if isinstance(UpperCAmelCase			,  UpperCAmelCase				):
              __snake_case				:			Any            =    (
                  LoRAAttnProcessor2_0 if hasattr(UpperCAmelCase			,  "scaled_dot_product_attention"				) else LoRAAttnProcessor
              )
              self.assertIsInstance(module.processor			,  UpperCAmelCase				)







      @unittest.skipIf(torch_device != "cuda"			,  "This test is supposed to run on GPU"				)
      def   UpperCAmelCase						(	self				)    ->				List[Any]:






        '''simple docstring'''






        with tempfile.TemporaryDirectory() as tmpdirname:
          self.create_lora_weight_file(UpperCAmelCase				)

          __snake_case				:			List[str]            =    self.get_dummy_components()
          __snake_case				:			str            =    StableDiffusionPipeline(**UpperCAmelCase				)
          __snake_case				:			int            =    sd_pipe.to(UpperCAmelCase				)
          sd_pipe.set_progress_bar_config(disable=UpperCAmelCase				)

          # enable XFormers
          sd_pipe.enable_xformers_memory_efficient_attention()

          # check if xFormers attention processors are used
          for _, module in sd_pipe.unet.named_modules():
            if isinstance(UpperCAmelCase			,  UpperCAmelCase				):
              self.assertIsInstance(module.processor			,  UpperCAmelCase				)

            # load LoRA weight file
          sd_pipe.load_lora_weights(UpperCAmelCase				)

          # check if lora attention processors are used
          for _, module in sd_pipe.unet.named_modules():
            if isinstance(UpperCAmelCase			,  UpperCAmelCase				):
              self.assertIsInstance(module.processor			,  UpperCAmelCase				)







      @unittest.skipIf(torch_device != "cuda"			,  "This test is supposed to run on GPU"				)
      def   UpperCAmelCase						(	self				)    ->				List[str]:






        '''simple docstring'''






        __snake_case				:			int            =    self.get_dummy_components()
        __snake_case				:			str            =    StableDiffusionPipeline(**UpperCAmelCase				)
        __snake_case				:			Dict            =    sd_pipe.to(UpperCAmelCase				)
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			str            =    self.get_dummy_inputs()

        # enable XFormers
        sd_pipe.enable_xformers_memory_efficient_attention()

        __snake_case				:			List[Any]            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			str            =    original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
          LoraLoaderMixin.save_lora_weights(
              save_directory=UpperCAmelCase			,  unet_lora_layers=lora_components["unet_lora_layers"]			,  text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,  )
          self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase			,  "pytorch_lora_weights.bin"				)				)				)
          sd_pipe.load_lora_weights(UpperCAmelCase				)

        __snake_case				:			Optional[int]            =    sd_pipe(**UpperCAmelCase				).images
        __snake_case				:			List[Any]            =    lora_images[0, -3:, -3:, -1]

        # Outputs shouldn't match.
        self.assertFalse(torch.allclose(torch.from_numpy(UpperCAmelCase				)			,  torch.from_numpy(UpperCAmelCase				)				)				)








from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


_UpperCamelCase    =	logging.get_logger(__name__)

_UpperCamelCase    =	{
    '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json''',
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}




class 			_lowerCamelCase	(      a      ):



      """simple docstring"""






      UpperCAmelCase_  : Tuple				      ="blenderbot"
      UpperCAmelCase_  : Any				      =["past_key_values"]
      UpperCAmelCase_  : Optional[int]				      ={"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}







      def __init__(	self			,  UpperCAmelCase=8008			,  UpperCAmelCase=128			,  UpperCAmelCase=2			,  UpperCAmelCase=10240			,  UpperCAmelCase=32			,  UpperCAmelCase=24			,  UpperCAmelCase=10240			,  UpperCAmelCase=32			,  UpperCAmelCase=0.0			,  UpperCAmelCase=0.0			,  UpperCAmelCase=True			,  UpperCAmelCase=True			,  UpperCAmelCase="gelu"			,  UpperCAmelCase=2560			,  UpperCAmelCase=0.1			,  UpperCAmelCase=0.0			,  UpperCAmelCase=0.0			,  UpperCAmelCase=0.02			,  UpperCAmelCase=1			,  UpperCAmelCase=False			,  UpperCAmelCase=0			,  UpperCAmelCase=1			,  UpperCAmelCase=2			,  UpperCAmelCase=3			,  UpperCAmelCase=2			,  **UpperCAmelCase			,  )    ->				List[Any]:






        '''simple docstring'''






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

        super().__init__(
            pad_token_id=UpperCAmelCase			,  bos_token_id=UpperCAmelCase			,  eos_token_id=UpperCAmelCase			,  is_encoder_decoder=UpperCAmelCase			,  decoder_start_token_id=UpperCAmelCase			,  encoder_no_repeat_ngram_size=UpperCAmelCase			,  forced_eos_token_id=UpperCAmelCase			,  **UpperCAmelCase			,  )




class 			_lowerCamelCase	(      a      ):



      """simple docstring"""






      @property
      def   UpperCAmelCase						(	self				)    ->				Mapping[str, Mapping[int, str]]:






        '''simple docstring'''






        if self.task in ["default", "seq2seq-lm"]:
          __snake_case				:			Any            =    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				:			Optional[Any]            =    {0: "batch", 1: "past_decoder_sequence + sequence"}
          else:
            __snake_case				:			int            =    {0: "batch", 1: "decoder_sequence"}
            __snake_case				:			Optional[int]            =    {0: "batch", 1: "decoder_sequence"}
          if self.use_past:
            self.fill_with_past_key_values_(UpperCAmelCase			,  direction="inputs"				)
        elif self.task == "causal-lm":
          __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				:			List[Any]            =    self.num_layers
            for i in range(UpperCAmelCase				):
              __snake_case				:			int            =    {0: "batch", 2: "past_sequence + sequence"}
              __snake_case				:			Tuple            =    {0: "batch", 2: "past_sequence + sequence"}
        else:
          __snake_case				:			Dict            =    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
      # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
      def   UpperCAmelCase						(	self				)    ->				Mapping[str, Mapping[int, str]]:






        '''simple docstring'''






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







      def   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase = -1			,  UpperCAmelCase = -1			,  UpperCAmelCase = False			,  UpperCAmelCase = None			,  )    ->				Mapping[str, Any]:






        '''simple docstring'''






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

          for _ in range(UpperCAmelCase				):
            common_inputs["past_key_values"].append(
                (
                    torch.zeros(UpperCAmelCase				),
                    torch.zeros(UpperCAmelCase				),
                    torch.zeros(UpperCAmelCase				),
                    torch.zeros(UpperCAmelCase				),
                )				)
        return common_inputs







      def   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase = -1			,  UpperCAmelCase = -1			,  UpperCAmelCase = False			,  UpperCAmelCase = None			,  )    ->				Mapping[str, Any]:






        '''simple docstring'''






        __snake_case				:			List[Any]            =    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				:			List[Any]            =    common_inputs["input_ids"].shape
          __snake_case				:			Any            =    seqlen
          __snake_case      ,				__snake_case				:			Dict            =    self.num_layers
          __snake_case      ,				__snake_case				:			int            =    self.num_attention_heads
          __snake_case				:			str            =    (
              batch,
              num_encoder_attention_heads,
              past_key_values_length,
              self._config.hidden_size // num_encoder_attention_heads,
          )
          __snake_case				:			Dict            =    common_inputs["attention_mask"].dtype
          __snake_case				:			Optional[Any]            =    torch.cat(
              [common_inputs["attention_mask"], torch.ones(UpperCAmelCase			,  UpperCAmelCase			,  dtype=UpperCAmelCase				)]			,  dim=1				)
          __snake_case				:			str            =    [
              (torch.zeros(UpperCAmelCase				), torch.zeros(UpperCAmelCase				)) for _ in range(UpperCAmelCase				)
          ]
        return common_inputs







      def   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase = -1			,  UpperCAmelCase = -1			,  UpperCAmelCase = False			,  UpperCAmelCase = None			,  )    ->				Mapping[str, Any]:






        '''simple docstring'''






        __snake_case				:			Tuple            =    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				:			str            =    tokenizer.num_special_tokens_to_add(UpperCAmelCase				)
        __snake_case				:			Any            =    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				:			Tuple            =    [" ".join([tokenizer.unk_token]				) * seq_length] * batch_size
        __snake_case				:			str            =    dict(tokenizer(UpperCAmelCase			,  return_tensors=UpperCAmelCase				)				)
        return common_inputs







      def   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase = -1			,  UpperCAmelCase = -1			,  UpperCAmelCase = False			,  UpperCAmelCase = None			,  )    ->				Mapping[str, Any]:






        '''simple docstring'''






        if self.task in ["default", "seq2seq-lm"]:
          __snake_case				:			Any            =    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				:			Union[str, Any]            =    self._generate_dummy_inputs_for_causal_lm(
              UpperCAmelCase			,  batch_size=UpperCAmelCase			,  seq_length=UpperCAmelCase			,  is_pair=UpperCAmelCase			,  framework=UpperCAmelCase				)
        else:
          __snake_case				:			Dict            =    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   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase				)    ->				int:






        '''simple docstring'''






        if self.task in ["default", "seq2seq-lm"]:
          __snake_case				:			str            =    super()._flatten_past_key_values_(UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase				)
        else:
          __snake_case				:			int            =    super(UpperCAmelCase			,  self				)._flatten_past_key_values_(
              UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase				)







      def   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase				)    ->				Any:






        '''simple docstring'''






        if direction not in ["inputs", "outputs"]:
          raise ValueError(F"""direction must either be \"inputs\" or \"outputs\", but {direction} was given"""				)

        __snake_case				:			Tuple            =    "past_key_values" if direction == "inputs" else "present"
        __snake_case      ,				__snake_case				:			Any            =    self.num_layers

        __snake_case				:			int            =    "past_encoder_sequence"
        __snake_case				:			int            =    "past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence"

        for i in range(UpperCAmelCase				):
          __snake_case				:			Any            =    {0: "batch", 2: decoder_sequence}
          __snake_case				:			List[Any]            =    {0: "batch", 2: decoder_sequence}
          __snake_case				:			Tuple            =    {0: "batch", 2: encoder_sequence}
          __snake_case				:			int            =    {0: "batch", 2: encoder_sequence}

