


"""simple docstring"""


import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
		    RwkvForCausalLM,
		    RwkvModel,
		)
		from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
		a_											=			False
class snake_case					:







		def __init__(    self					:				Union[str, Any]			,      a__					:				Optional[int]			,      a__					:				Any=14			,      a__					:				Union[str, Any]=7			,      a__					:				List[str]=True			,      a__					:				Optional[Any]=False			,      a__					:				str=True			,      a__					:				Optional[int]=True			,      a__					:				Any=True			,      a__					:				Optional[int]=99			,      a__					:				Dict=32			,      a__					:				Any=5			,      a__					:				Any=37			,      a__					:				Tuple="gelu"			,      a__					:				Optional[int]=0.1			,      a__					:				Optional[Any]=0.1			,      a__					:				Dict=5_12			,      a__					:				int=16			,      a__					:				str=2			,      a__					:				List[Any]=3			,      a__					:				Dict=4			,      a__					:				Union[str, Any]=None			,      )				->  int:





						'''simple docstring'''







						_A         =	parent
						_A         =	batch_size
						_A         =	seq_length
						_A         =	is_training
						_A         =	use_token_type_ids
						_A         =	use_input_mask
						_A         =	use_labels
						_A         =	use_mc_token_ids
						_A         =	vocab_size
						_A         =	hidden_size
						_A         =	num_hidden_layers
						_A         =	intermediate_size
						_A         =	hidden_act
						_A         =	hidden_dropout_prob
						_A         =	attention_probs_dropout_prob
						_A         =	max_position_embeddings
						_A         =	type_vocab_size
						_A         =	type_sequence_label_size
						_A         =	num_labels
						_A         =	num_choices
						_A         =	scope
						_A         =	vocab_size - 1
						_A         =	vocab_size - 1
						_A         =	vocab_size - 1







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





						'''simple docstring'''







						return RwkvConfig.from_pretrained("sgugger/rwkv-4-pile-7b"						)







		def        a_       (    self					:				Tuple			,      a__					:				Union[str, Any]=False			,      a__					:				Any=False			,      a__					:				Union[str, Any]=False						)				->  Tuple:





						'''simple docstring'''







						_A         =	ids_tensor([self.batch_size, self.seq_length]			,      self.vocab_size						)

						_A         =	None
						if self.use_input_mask:
										_A         =	random_attention_mask([self.batch_size, self.seq_length]						)

						_A         =	None
						if self.use_token_type_ids:
										_A         =	ids_tensor([self.batch_size, self.seq_length]			,      self.type_vocab_size						)

						_A         =	None
						if self.use_mc_token_ids:
										_A         =	ids_tensor([self.batch_size, self.num_choices]			,      self.seq_length						)

						_A         =	None
						_A         =	None
						_A         =	None
						if self.use_labels:
										_A         =	ids_tensor([self.batch_size]			,      self.type_sequence_label_size						)
										_A         =	ids_tensor([self.batch_size, self.seq_length]			,      self.num_labels						)
										_A         =	ids_tensor([self.batch_size]			,      self.num_choices						)

						_A         =	self.get_config(
						    gradient_checkpointing=a__			,      scale_attn_by_inverse_layer_idx=a__			,      reorder_and_upcast_attn=a__			,      )

						return (
						    config,
						    input_ids,
						    input_mask,
						    None,
						    token_type_ids,
						    mc_token_ids,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						)







		def        a_       (    self					:				str			,      a__					:				str=False			,      a__					:				Union[str, Any]=False			,      a__					:				str=False						)				->  Optional[Any]:





						'''simple docstring'''







						return RwkvConfig(
						    vocab_size=self.vocab_size			,      hidden_size=self.hidden_size			,      num_hidden_layers=self.num_hidden_layers			,      intermediate_size=self.intermediate_size			,      activation_function=self.hidden_act			,      resid_pdrop=self.hidden_dropout_prob			,      attn_pdrop=self.attention_probs_dropout_prob			,      n_positions=self.max_position_embeddings			,      type_vocab_size=self.type_vocab_size			,      use_cache=a__			,      bos_token_id=self.bos_token_id			,      eos_token_id=self.eos_token_id			,      pad_token_id=self.pad_token_id			,      gradient_checkpointing=a__			,      scale_attn_by_inverse_layer_idx=a__			,      reorder_and_upcast_attn=a__			,      )







		def        a_       (    self					:				List[str]						)				->  Optional[Any]:





						'''simple docstring'''







						_A         =	self.get_config()
						_A         =	3_00
						return config







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





						'''simple docstring'''







						(
						    (
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						
						)         =	self.prepare_config_and_inputs()

						_A         =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
						_A         =	ids_tensor([self.batch_size, self.seq_length]			,      vocab_size=2						)

						return (
						    config,
						    input_ids,
						    input_mask,
						    head_mask,
						    token_type_ids,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)







		def        a_       (    self					:				Optional[Any]			,      a__					:				List[Any]			,      a__					:				int			,      a__					:				List[str]			,      a__					:				List[str]			,      a__					:				str			,      *a__					:				List[Any]						)				->  Any:





						'''simple docstring'''







						_A         =	True
						_A         =	RwkvModel(config=a__						)
						model.to(a__						)
						model.eval()

						_A         =	model(a__						)

						self.parent.assertEqual(result.last_hidden_state.shape			,      (self.batch_size, self.seq_length, self.hidden_size)						)
						self.parent.assertEqual(len(result.hidden_states						)			,      config.num_hidden_layers + 1						)







		def        a_       (    self					:				Optional[Any]			,      a__					:				Dict			,      a__					:				Optional[Any]			,      a__					:				List[Any]			,      a__					:				List[Any]			,      a__					:				Optional[int]			,      *a__					:				Optional[int]						)				->  Optional[int]:





						'''simple docstring'''







						_A         =	RwkvForCausalLM(a__						)
						model.to(a__						)
						model.eval()

						_A         =	model(a__			,      labels=a__						)
						self.parent.assertEqual(result.loss.shape			,      ()						)
						self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.seq_length, self.vocab_size)						)







		def        a_       (    self					:				Dict			,      a__					:				List[str]			,      a__					:				Any			,      a__					:				Optional[int]			,      a__					:				List[str]			,      a__					:				Any			,      *a__					:				Optional[Any]						)				->  Dict:





						'''simple docstring'''







						_A         =	RwkvModel(config=a__						)
						model.to(a__						)
						model.eval()

						_A         =	model(a__						)
						_A         =	outputs.last_hidden_state

						_A         =	model(input_ids[:, :2]						)
						_A         =	outputs.last_hidden_state

						# Using the state computed on the first inputs, we will get the same output
						_A         =	model(input_ids[:, 2:]			,      state=outputs.state						)
						_A         =	outputs.last_hidden_state

						self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]			,      dim=1						)			,      a__			,      atol=1E-5						)						)







		def        a_       (    self					:				Dict			,      a__					:				Union[str, Any]			,      a__					:				Dict			,      a__					:				Optional[int]			,      a__					:				Any			,      a__					:				List[str]			,      *a__					:				Tuple			,      a__					:				Optional[Any]=False						)				->  Any:





						'''simple docstring'''







						_A         =	RwkvForCausalLM(a__						)
						model.to(a__						)
						if gradient_checkpointing:
										model.gradient_checkpointing_enable()

						_A         =	model(a__			,      labels=a__						)
						self.parent.assertEqual(result.loss.shape			,      ()						)
						self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.seq_length, self.vocab_size)						)
						result.loss.backward()







		def        a_       (    self					:				Union[str, Any]						)				->  Optional[int]:





						'''simple docstring'''







						_A         =	self.prepare_config_and_inputs()

						(
						    (
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						(
						    _A
						)     ,						
						)         =	config_and_inputs

						_A         =	{"input_ids": input_ids}

						return config, inputs_dict





@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0   , reason='See https://github.com/huggingface/transformers/pull/24204')
@require_torch
class snake_case					(     _UpperCamelCase   , _UpperCamelCase   , _UpperCamelCase   , unittest.TestCase):
		__UpperCamelCase							=	(RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
		__UpperCamelCase							=	(
		    {'feature-extraction': RwkvModel, 'text-generation': RwkvForCausalLM} if is_torch_available() else {}
		)
		# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
		__UpperCamelCase							=	False
		__UpperCamelCase							=	False
		__UpperCamelCase							=	False
		__UpperCamelCase							=	False
		__UpperCamelCase							=	False  # Rwkv does not support head masking







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





						'''simple docstring'''







						_A         =	RwkvModelTester(self						)
						_A         =	ConfigTester(
						    self			,      config_class=a__			,      n_embd=37			,      common_properties=["hidden_size", "num_hidden_layers"]						)







		def        a_       (    self					:				Dict			,      a__					:				int			,      a__					:				Dict			,      a__					:				Optional[Any]=None						)				->  Optional[Any]:





						'''simple docstring'''







						if isinstance(a__			,      torch.Tensor						):
										_A     ,						_A         =	member.max().item(), member.min().item()
						elif isinstance(a__			,      a__						) or isinstance(a__			,      a__						):
										_A     ,						_A         =	max(a__						), min(a__						)

						if not isinstance(a__			,      a__						):
										raise TypeError("container should be a list or tuple"						)
						elif len(a__						) != 2:
										raise ValueError("container should have 2 elements"						)

						_A     ,						_A         =	container

						_A         =	(min_value >= expected_min) and (max_value <= expected_max)

						if not is_inside_interval:
										_A         =	"%s not found in %s" % (safe_repr(a__						), safe_repr(a__						))
										self.fail(self._formatMessage(a__			,      a__						)						)







		def        a_       (    self					:				Tuple						)				->  int:





						'''simple docstring'''







						self.config_tester.run_common_tests()







		def        a_       (    self					:				Optional[Any]						)				->  List[Any]:





						'''simple docstring'''







						_A         =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_rwkv_model(*a__						)







		def        a_       (    self					:				Optional[int]						)				->  int:





						'''simple docstring'''







						_A         =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_causl_lm(*a__						)







		def        a_       (    self					:				str						)				->  Any:





						'''simple docstring'''







						_A         =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_state_equivalency(*a__						)







		def        a_       (    self					:				Any						)				->  int:





						'''simple docstring'''







						_A     ,						_A         =	self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										_A         =	model_class(config=a__						)
										for name, param in model.named_parameters():
														if "time_decay" in name:
																		if param.requires_grad:
																						self.assertTrue(param.data.max().item() == 3.0						)
																						self.assertTrue(param.data.min().item() == -5.0						)
														elif "time_first" in name:
																		if param.requires_grad:
																						# check if it's a ones like
																						self.assertTrue(torch.allclose(param.data			,      torch.ones_like(param.data						)			,      atol=1E-5			,      rtol=1E-5						)						)
														elif any(x in name for x in ["time_mix_key", "time_mix_receptance"]						):
																		if param.requires_grad:
																						self.assertInterval(
																						    param.data			,      [0.0, 1.0]			,      msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""			,      )
														elif "time_mix_value" in name:
																		if param.requires_grad:
																						self.assertInterval(
																						    param.data			,      [0.0, 1.3]			,      msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""			,      )







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





						'''simple docstring'''







						_A     ,						_A         =	self.model_tester.prepare_config_and_inputs_for_common()
						_A         =	True

						_A         =	getattr(self.model_tester			,      "seq_length"			,      a__						)

						for model_class in self.all_model_classes:
										_A         =	True
										_A         =	False
										_A         =	True
										_A         =	model_class(a__						)
										model.to(a__						)
										model.eval()

										_A         =	self._prepare_for_class(a__			,      a__						)
										_A         =	inputs["input_ids"].shape[0]
										with torch.no_grad():
														_A         =	model(**a__						)
										_A         =	outputs.attentions
										self.assertEqual(len(a__						)			,      self.model_tester.num_hidden_layers						)

										# check that output_attentions also work using config
										del inputs_dict["output_attentions"]
										_A         =	True
										_A         =	model_class(a__						)
										model.to(a__						)
										model.eval()

										_A         =	self._prepare_for_class(a__			,      a__						)
										_A         =	inputs["input_ids"].shape[0]
										with torch.no_grad():
														_A         =	model(**a__						)
										_A         =	outputs.attentions
										self.assertEqual(len(a__						)			,      self.model_tester.num_hidden_layers						)

										self.assertListEqual(
										    list(attentions[0].shape[-3:]						)			,      [batch_size, seq_len, config.hidden_size]			,      )
										_A         =	len(a__						)

										# Check attention is always last and order is fine
										_A         =	True
										_A         =	True
										_A         =	model_class(a__						)
										model.to(a__						)
										model.eval()

										_A         =	self._prepare_for_class(a__			,      a__						)
										_A         =	inputs["input_ids"].shape[0]
										with torch.no_grad():
														_A         =	model(**a__						)

										_A         =	1
										self.assertEqual(out_len + added_hidden_states			,      len(a__						)						)

										_A         =	outputs.attentions

										self.assertEqual(len(a__						)			,      self.model_tester.num_hidden_layers						)
										self.assertListEqual(
										    list(self_attentions[0].shape[-3:]						)			,      [batch_size, seq_len, config.hidden_size]			,      )







		@slow
		def        a_       (    self					:				str						)				->  Dict:





						'''simple docstring'''







						for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										_A         =	RwkvModel.from_pretrained(a__						)
										self.assertIsNotNone(a__						)





@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0   , reason='See https://github.com/huggingface/transformers/pull/24204')
@slow
class snake_case					(     unittest.TestCase):







		def        a_       (    self					:				int						)				->  Dict:





						'''simple docstring'''







						_A         =	"RWKV/rwkv-4-169m-pile"
						_A         =	AutoTokenizer.from_pretrained(self.model_id						)







		def        a_       (    self					:				List[Any]						)				->  Dict:





						'''simple docstring'''







						_A         =	"Hello my name is Jasmine and I am a newbie to the"
						_A         =	RwkvForCausalLM.from_pretrained(self.model_id						).to(a__						)

						_A         =	self.tokenizer("Hello my name is"			,      return_tensors="pt"						).input_ids.to(a__						)
						_A         =	model.generate(a__			,      max_new_tokens=10						)
						_A         =	self.tokenizer.decode(output[0].tolist()						)

						self.assertEqual(a__			,      a__						)







		def        a_       (    self					:				Optional[int]						)				->  Optional[int]:





						'''simple docstring'''







						_A         =	"Hello my name is Jasmine and I am a newbie to the"

						_A         =	self.tokenizer("Hello my name is"			,      return_tensors="pt"						).input_ids.to(a__						)
						_A         =	RwkvForCausalLM.from_pretrained(self.model_id			,      torch_dtype=torch.bfloataa						).to(a__						)

						_A         =	model.generate(a__			,      max_new_tokens=10						)
						_A         =	self.tokenizer.decode(output[0].tolist()						)

						self.assertEqual(a__			,      a__						)




'''simple docstring'''

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

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


UpperCAmelCase			=    logging.get_logger(__name__)

UpperCAmelCase			=    {
    '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 __snake_case(      _lowerCAmelCase					):


						'''simple docstring'''



						UpperCAmelCase							:    Union[str, Any]           =		"bigbird_pegasus"
						UpperCAmelCase							:    Optional[Any]           =		["past_key_values"]
						UpperCAmelCase							:    Optional[Any]           =		{
						    "num_attention_heads": "encoder_attention_heads",
						    "hidden_size": "d_model",
						    "attention_probs_dropout_prob": "attention_dropout",
						}


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

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

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







class __snake_case(      _lowerCAmelCase					):


						'''simple docstring'''



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

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

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


						def 				__snake_case    (					self	,       A_	,       A_ = -1	,       A_ = -1	,       A_ = False	,       A_ = None	,       )	->	Mapping[str, Any]:
							lowerCAmelCase    =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    A_	,       A_	,       A_	,       A_	,       A_      )

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

							if self.use_past:
								if not is_torch_available():
									raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""      )
								else:
									import torch
								lowerCAmelCase,   lowerCAmelCase    =	common_inputs["""input_ids"""].shape
								lowerCAmelCase    =	common_inputs["""decoder_input_ids"""].shape[1]
								lowerCAmelCase,   lowerCAmelCase    =	self.num_attention_heads
								lowerCAmelCase    =	(
								    batch,
								    num_encoder_attention_heads,
								    encoder_seq_length,
								    self._config.hidden_size // num_encoder_attention_heads,
								)
								lowerCAmelCase    =	decoder_seq_length + 3
								lowerCAmelCase    =	(
								    batch,
								    num_decoder_attention_heads,
								    decoder_past_length,
								    self._config.hidden_size // num_decoder_attention_heads,
								)

								lowerCAmelCase    =	torch.cat(
								    [common_inputs["""decoder_attention_mask"""], torch.ones(A_	,       A_      )]	,       dim=1      )

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

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


						def 				__snake_case    (					self	,       A_	,       A_ = -1	,       A_ = -1	,       A_ = False	,       A_ = None	,       )	->	Mapping[str, Any]:
							lowerCAmelCase    =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    A_	,       A_	,       A_	,       A_	,       A_      )

							if self.use_past:
								if not is_torch_available():
									raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""      )
								else:
									import torch
								lowerCAmelCase,   lowerCAmelCase    =	common_inputs["""input_ids"""].shape
								# Not using the same length for past_key_values
								lowerCAmelCase    =	seqlen + 2
								lowerCAmelCase,   lowerCAmelCase    =	self.num_layers
								lowerCAmelCase,   lowerCAmelCase    =	self.num_attention_heads
								lowerCAmelCase    =	(
								    batch,
								    num_encoder_attention_heads,
								    past_key_values_length,
								    self._config.hidden_size // num_encoder_attention_heads,
								)

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


						def 				__snake_case    (					self	,       A_	,       A_ = -1	,       A_ = -1	,       A_ = False	,       A_ = 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
							lowerCAmelCase    =	compute_effective_axis_dimension(
							    A_	,       fixed_dimension=OnnxConfig.default_fixed_batch	,       num_token_to_add=0      )

							# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
							lowerCAmelCase    =	tokenizer.num_special_tokens_to_add(A_      )
							lowerCAmelCase    =	compute_effective_axis_dimension(
							    A_	,       fixed_dimension=OnnxConfig.default_fixed_sequence	,       num_token_to_add=A_      )

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


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

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

							return common_inputs




						def 				__snake_case    (					self	,       A_	,       A_	,       A_	,       A_      )	->	Tuple:
							if self.task in ["default", "seq2seq-lm"]:
								lowerCAmelCase    =	super()._flatten_past_key_values_(A_	,       A_	,       A_	,       A_      )
							else:
								lowerCAmelCase    =	super(A_	,       self      )._flatten_past_key_values_(
								    A_	,       A_	,       A_	,       A_      )