





"""simple docstring"""





import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor




def 		lowerCAmelCase_	(  lowercase_   :							Union[str, Any]		):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE :			Optional[Any]												=       TimesformerConfig()

 if "large" in model_name:
  __SCREAMING_SNAKE_CASE :			Any												=       96

 if "hr" in model_name:
  __SCREAMING_SNAKE_CASE :			List[str]												=       16
  __SCREAMING_SNAKE_CASE :			Tuple												=       448

 __SCREAMING_SNAKE_CASE :			Tuple												=       '''huggingface/label-files'''
 if "k400" in model_name:
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       400
  __SCREAMING_SNAKE_CASE :			Tuple												=       '''kinetics400-id2label.json'''
 elif "k600" in model_name:
  __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       600
  __SCREAMING_SNAKE_CASE :			List[str]												=       '''kinetics600-id2label.json'''
 elif "ssv2" in model_name:
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       174
  __SCREAMING_SNAKE_CASE :			Tuple												=       '''something-something-v2-id2label.json'''
 else:
  raise ValueError('''Model name should either contain \'k400\', \'k600\' or \'ssv2\'.'''		)
 __SCREAMING_SNAKE_CASE :			str												=       json.load(open(hf_hub_download(lowercase_ ,      lowercase_ ,      repo_type='''dataset'''		) ,      '''r'''		)		)
 __SCREAMING_SNAKE_CASE :			Any												=       {int(lowercase_		): v for k, v in idalabel.items()}
 __SCREAMING_SNAKE_CASE :			Any												=       idalabel
 __SCREAMING_SNAKE_CASE :			Optional[Any]												=       {v: k for k, v in idalabel.items()}

 return config




def 		lowerCAmelCase_	(  lowercase_   :							Optional[int]		):

 '''simple docstring'''






 if "encoder." in name:
  __SCREAMING_SNAKE_CASE :			List[Any]												=       name.replace('''encoder.''' ,      ''''''		)
 if "cls_token" in name:
  __SCREAMING_SNAKE_CASE :			List[str]												=       name.replace('''cls_token''' ,      '''timesformer.embeddings.cls_token'''		)
 if "pos_embed" in name:
  __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       name.replace('''pos_embed''' ,      '''timesformer.embeddings.position_embeddings'''		)
 if "time_embed" in name:
  __SCREAMING_SNAKE_CASE :			List[str]												=       name.replace('''time_embed''' ,      '''timesformer.embeddings.time_embeddings'''		)
 if "patch_embed.proj" in name:
  __SCREAMING_SNAKE_CASE :			List[Any]												=       name.replace('''patch_embed.proj''' ,      '''timesformer.embeddings.patch_embeddings.projection'''		)
 if "patch_embed.norm" in name:
  __SCREAMING_SNAKE_CASE :			str												=       name.replace('''patch_embed.norm''' ,      '''timesformer.embeddings.norm'''		)
 if "blocks" in name:
  __SCREAMING_SNAKE_CASE :			Any												=       name.replace('''blocks''' ,      '''timesformer.encoder.layer'''		)
 if "attn.proj" in name:
  __SCREAMING_SNAKE_CASE :			str												=       name.replace('''attn.proj''' ,      '''attention.output.dense'''		)
 if "attn" in name and "bias" not in name and "temporal" not in name:
  __SCREAMING_SNAKE_CASE :			Tuple												=       name.replace('''attn''' ,      '''attention.self'''		)
 if "attn" in name and "temporal" not in name:
  __SCREAMING_SNAKE_CASE :			Dict												=       name.replace('''attn''' ,      '''attention.attention'''		)
 if "temporal_norm1" in name:
  __SCREAMING_SNAKE_CASE :			Dict												=       name.replace('''temporal_norm1''' ,      '''temporal_layernorm'''		)
 if "temporal_attn.proj" in name:
  __SCREAMING_SNAKE_CASE :			int												=       name.replace('''temporal_attn''' ,      '''temporal_attention.output.dense'''		)
 if "temporal_fc" in name:
  __SCREAMING_SNAKE_CASE :			List[Any]												=       name.replace('''temporal_fc''' ,      '''temporal_dense'''		)
 if "norm1" in name and "temporal" not in name:
  __SCREAMING_SNAKE_CASE :			Tuple												=       name.replace('''norm1''' ,      '''layernorm_before'''		)
 if "norm2" in name:
  __SCREAMING_SNAKE_CASE :			List[str]												=       name.replace('''norm2''' ,      '''layernorm_after'''		)
 if "mlp.fc1" in name:
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       name.replace('''mlp.fc1''' ,      '''intermediate.dense'''		)
 if "mlp.fc2" in name:
  __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       name.replace('''mlp.fc2''' ,      '''output.dense'''		)
 if "norm.weight" in name and "fc" not in name and "temporal" not in name:
  __SCREAMING_SNAKE_CASE :			List[Any]												=       name.replace('''norm.weight''' ,      '''timesformer.layernorm.weight'''		)
 if "norm.bias" in name and "fc" not in name and "temporal" not in name:
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       name.replace('''norm.bias''' ,      '''timesformer.layernorm.bias'''		)
 if "head" in name:
  __SCREAMING_SNAKE_CASE :			int												=       name.replace('''head''' ,      '''classifier'''		)

 return name




def 		lowerCAmelCase_	(  lowercase_   :							List[str] ,      lowercase_   :							Tuple		):

 '''simple docstring'''






 for key in orig_state_dict.copy().keys():
  __SCREAMING_SNAKE_CASE :			int												=       orig_state_dict.pop(lowercase_		)

  if key.startswith('''model.'''		):
   __SCREAMING_SNAKE_CASE :			Dict												=       key.replace('''model.''' ,      ''''''		)

  if "qkv" in key:
   __SCREAMING_SNAKE_CASE :			str												=       key.split('''.'''		)
   __SCREAMING_SNAKE_CASE :			str												=       int(key_split[1]		)
   __SCREAMING_SNAKE_CASE :			Optional[Any]												=       '''timesformer.encoder.layer.'''
   if "temporal" in key:
    __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       '''.temporal_attention.attention.qkv.'''
   else:
    __SCREAMING_SNAKE_CASE :			Any												=       '''.attention.attention.qkv.'''
   if "weight" in key:
    __SCREAMING_SNAKE_CASE :			Optional[int]												=       val
   else:
    __SCREAMING_SNAKE_CASE :			List[Any]												=       val
  else:
   __SCREAMING_SNAKE_CASE :			Tuple												=       val

 return orig_state_dict




def 		lowerCAmelCase_	(  ):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE :			List[Any]												=       hf_hub_download(
     repo_id='''hf-internal-testing/spaghetti-video''' ,      filename='''eating_spaghetti.npy''' ,      repo_type='''dataset'''		)
 __SCREAMING_SNAKE_CASE :			str												=       np.load(lowercase_		)
 return list(lowercase_		)




def 		lowerCAmelCase_	(  lowercase_   :							Any ,      lowercase_   :							Tuple ,      lowercase_   :							Dict ,      lowercase_   :							Any		):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       get_timesformer_config(lowercase_		)

 __SCREAMING_SNAKE_CASE :			str												=       TimesformerForVideoClassification(lowercase_		)

 # download original checkpoint, hosted on Google Drive
 __SCREAMING_SNAKE_CASE :			Any												=       '''pytorch_model.bin'''
 gdown.cached_download(lowercase_ ,      lowercase_ ,      quiet=lowercase_		)
 __SCREAMING_SNAKE_CASE :			Any												=       torch.load(lowercase_ ,      map_location='''cpu'''		)
 if "model" in files:
  __SCREAMING_SNAKE_CASE :			Any												=       files['''model''']
 elif "module" in files:
  __SCREAMING_SNAKE_CASE :			str												=       files['''module''']
 else:
  __SCREAMING_SNAKE_CASE :			str												=       files['''model_state''']
 __SCREAMING_SNAKE_CASE :			int												=       convert_state_dict(lowercase_ ,      lowercase_		)

 model.load_state_dict(lowercase_		)
 model.eval()

 # verify model on basic input
 __SCREAMING_SNAKE_CASE :			Optional[Any]												=       VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5] ,      image_std=[0.5, 0.5, 0.5]		)
 __SCREAMING_SNAKE_CASE :			str												=       prepare_video()
 __SCREAMING_SNAKE_CASE :			Any												=       image_processor(video[:8] ,      return_tensors='''pt'''		)

 __SCREAMING_SNAKE_CASE :			List[Any]												=       model(**lowercase_		)
 __SCREAMING_SNAKE_CASE :			str												=       outputs.logits

 __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       [
     # Kinetics-400 checkpoints (hr = high resolution input of 448px instead of 224px)
     '''timesformer-base-finetuned-k400''',
     '''timesformer-large-finetuned-k400''',
     '''timesformer-hr-finetuned-k400''',
     # Kinetics-600 checkpoints (hr = high resolution input of 448px instead of 224px)
     '''timesformer-base-finetuned-k600''',
     '''timesformer-large-finetuned-k600''',
     '''timesformer-hr-finetuned-k600''',
     # Something-Something-v2 checkpoints (hr = high resolution input of 448px instead of 224px)
     '''timesformer-base-finetuned-ssv2''',
     '''timesformer-large-finetuned-ssv2''',
     '''timesformer-hr-finetuned-ssv2''',
 ]

 # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
 if model_name == "timesformer-base-finetuned-k400":
  __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       torch.Size([1, 400]		)
  __SCREAMING_SNAKE_CASE :			int												=       torch.tensor([-0.3_016, -0.7_713, -0.4_205]		)
 elif model_name == "timesformer-base-finetuned-k600":
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       torch.Size([1, 600]		)
  __SCREAMING_SNAKE_CASE :			Dict												=       torch.tensor([-0.7_267, -0.7_466, 3.2_404]		)
 elif model_name == "timesformer-base-finetuned-ssv2":
  __SCREAMING_SNAKE_CASE :			str												=       torch.Size([1, 174]		)
  __SCREAMING_SNAKE_CASE :			List[str]												=       torch.tensor([-0.9_059, 0.6_433, -3.1_457]		)
 elif model_name == "timesformer-large-finetuned-k400":
  __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       torch.Size([1, 400]		)
  __SCREAMING_SNAKE_CASE :			List[Any]												=       torch.tensor([0, 0, 0]		)
 elif model_name == "timesformer-large-finetuned-k600":
  __SCREAMING_SNAKE_CASE :			Dict												=       torch.Size([1, 600]		)
  __SCREAMING_SNAKE_CASE :			List[Any]												=       torch.tensor([0, 0, 0]		)
 elif model_name == "timesformer-large-finetuned-ssv2":
  __SCREAMING_SNAKE_CASE :			int												=       torch.Size([1, 174]		)
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       torch.tensor([0, 0, 0]		)
 elif model_name == "timesformer-hr-finetuned-k400":
  __SCREAMING_SNAKE_CASE :			Any												=       torch.Size([1, 400]		)
  __SCREAMING_SNAKE_CASE :			Tuple												=       torch.tensor([-0.9_617, -3.7_311, -3.7_708]		)
 elif model_name == "timesformer-hr-finetuned-k600":
  __SCREAMING_SNAKE_CASE :			Any												=       torch.Size([1, 600]		)
  __SCREAMING_SNAKE_CASE :			Tuple												=       torch.tensor([2.5_273, 0.7_127, 1.8_848]		)
 elif model_name == "timesformer-hr-finetuned-ssv2":
  __SCREAMING_SNAKE_CASE :			List[Any]												=       torch.Size([1, 174]		)
  __SCREAMING_SNAKE_CASE :			Optional[int]												=       torch.tensor([-3.6_756, -0.7_513, 0.7_180]		)
 else:
  raise ValueError(F'''Model name not supported. Should be one of {model_names}'''		)

 # verify logits
 assert logits.shape == expected_shape
 assert torch.allclose(logits[0, :3] ,      lowercase_ ,      atol=1E-4		)
 print('''Logits ok!'''		)

 if pytorch_dump_folder_path is not None:
  print(F'''Saving model and image processor to {pytorch_dump_folder_path}'''		)
  image_processor.save_pretrained(lowercase_		)
  model.save_pretrained(lowercase_		)

 if push_to_hub:
  print('''Pushing to the hub...'''		)
  model.push_to_hub(F'''fcakyon/{model_name}'''		)


if __name__ == "__main__":
    _lowerCamelCase       =	argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '''--checkpoint_url''',
        default='''https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download''',
        type=str,
        help=(
            '''URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'''
            ''' download link.'''
        ),
    )
    parser.add_argument(
        '''--pytorch_dump_folder_path''',
        default='''''',
        type=str,
        help='''Path to the output PyTorch model directory.''',
    )
    parser.add_argument('''--model_name''', default='''timesformer-base-finetuned-k400''', type=str, help='''Name of the model.''')
    parser.add_argument(
        '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
    )

    _lowerCamelCase       =	parser.parse_args()
    convert_timesformer_checkpoint(
        args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
    )









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 )
