

"""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:
						lowerCamelCase__			:					Optional[int]							  =					False






class       _UpperCAmelCase     :
			def __init__(						self ,						_A ,						_A=14 ,						_A=7 ,						_A=True ,						_A=False ,						_A=True ,						_A=True ,						_A=True ,						_A=99 ,						_A=32 ,						_A=5 ,						_A=37 ,						_A="gelu" ,						_A=0.1 ,						_A=0.1 ,						_A=5_12 ,						_A=16 ,						_A=2 ,						_A=3 ,						_A=4 ,						_A=None ,						)      ->      List[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Dict       =   parent
				_UpperCAmelCase      :							Optional[Any]       =   batch_size
				_UpperCAmelCase      :							Optional[Any]       =   seq_length
				_UpperCAmelCase      :							List[str]       =   is_training
				_UpperCAmelCase      :							Tuple       =   use_token_type_ids
				_UpperCAmelCase      :							Dict       =   use_input_mask
				_UpperCAmelCase      :							Any       =   use_labels
				_UpperCAmelCase      :							Dict       =   use_mc_token_ids
				_UpperCAmelCase      :							Tuple       =   vocab_size
				_UpperCAmelCase      :							Optional[Any]       =   hidden_size
				_UpperCAmelCase      :							Dict       =   num_hidden_layers
				_UpperCAmelCase      :							int       =   intermediate_size
				_UpperCAmelCase      :							Any       =   hidden_act
				_UpperCAmelCase      :							Optional[Any]       =   hidden_dropout_prob
				_UpperCAmelCase      :							int       =   attention_probs_dropout_prob
				_UpperCAmelCase      :							List[Any]       =   max_position_embeddings
				_UpperCAmelCase      :							str       =   type_vocab_size
				_UpperCAmelCase      :							Tuple       =   type_sequence_label_size
				_UpperCAmelCase      :							Optional[int]       =   num_labels
				_UpperCAmelCase      :							Union[str, Any]       =   num_choices
				_UpperCAmelCase      :							Tuple       =   scope
				_UpperCAmelCase      :							str       =   vocab_size - 1
				_UpperCAmelCase      :							str       =   vocab_size - 1
				_UpperCAmelCase      :							Optional[Any]       =   vocab_size - 1
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				return RwkvConfig.from_pretrained("""sgugger/rwkv-4-pile-7b"""       )
			def 							__snake_case	(						self ,						_A=False ,						_A=False ,						_A=False       )      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[int]       =   ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size       )

				_UpperCAmelCase      :							Union[str, Any]       =   None
				if self.use_input_mask:
					_UpperCAmelCase      :							Optional[Any]       =   random_attention_mask([self.batch_size, self.seq_length]       )

				_UpperCAmelCase      :							List[Any]       =   None
				if self.use_token_type_ids:
					_UpperCAmelCase      :							str       =   ids_tensor([self.batch_size, self.seq_length] ,						self.type_vocab_size       )

				_UpperCAmelCase      :							Optional[Any]       =   None
				if self.use_mc_token_ids:
					_UpperCAmelCase      :							List[Any]       =   ids_tensor([self.batch_size, self.num_choices] ,						self.seq_length       )

				_UpperCAmelCase      :							Optional[int]       =   None
				_UpperCAmelCase      :							Optional[int]       =   None
				_UpperCAmelCase      :							str       =   None
				if self.use_labels:
					_UpperCAmelCase      :							Tuple       =   ids_tensor([self.batch_size] ,						self.type_sequence_label_size       )
					_UpperCAmelCase      :							Union[str, Any]       =   ids_tensor([self.batch_size, self.seq_length] ,						self.num_labels       )
					_UpperCAmelCase      :							int       =   ids_tensor([self.batch_size] ,						self.num_choices       )

				_UpperCAmelCase      :							str       =   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 							__snake_case	(						self ,						_A=False ,						_A=False ,						_A=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 							__snake_case	(						self       )      ->      Any:






				'''simple docstring'''







				_UpperCAmelCase      :							Tuple       =   self.get_config()
				_UpperCAmelCase      :							Any       =   3_00
				return config
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				(
				    (
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				
				)      :							Optional[Any]       =   self.prepare_config_and_inputs()

				_UpperCAmelCase      :							List[str]       =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
				_UpperCAmelCase      :							Any       =   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 							__snake_case	(						self ,						_A ,						_A ,						_A ,						_A ,						_A ,						*_A       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							Any       =   True
				_UpperCAmelCase      :							List[str]       =   RwkvModel(config=_A       )
				model.to(_A       )
				model.eval()

				_UpperCAmelCase      :							List[Any]       =   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 							__snake_case	(						self ,						_A ,						_A ,						_A ,						_A ,						_A ,						*_A       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							Tuple       =   RwkvForCausalLM(_A       )
				model.to(_A       )
				model.eval()

				_UpperCAmelCase      :							str       =   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 							__snake_case	(						self ,						_A ,						_A ,						_A ,						_A ,						_A ,						*_A       )      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							List[Any]       =   RwkvModel(config=_A       )
				model.to(_A       )
				model.eval()

				_UpperCAmelCase      :							Optional[int]       =   model(_A       )
				_UpperCAmelCase      :							List[Any]       =   outputs.last_hidden_state

				_UpperCAmelCase      :							Union[str, Any]       =   model(input_ids[:, :2]       )
				_UpperCAmelCase      :							List[Any]       =   outputs.last_hidden_state

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

				self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two] ,						dim=1       ) ,						_A ,						atol=1e-5       )       )
			def 							__snake_case	(						self ,						_A ,						_A ,						_A ,						_A ,						_A ,						*_A ,						_A=False       )      ->      Optional[Any]:






				'''simple docstring'''







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

				_UpperCAmelCase      :							Union[str, Any]       =   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 							__snake_case	(						self       )      ->      Optional[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   self.prepare_config_and_inputs()

				(
				    (
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				(
				    _UpperCAmelCase
				)  ,				
				)      :							str       =   config_and_inputs

				_UpperCAmelCase      :							Union[str, Any]       =   {"""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       _UpperCAmelCase     (						__a      ,       __a      ,       __a      ,       unittest.TestCase):
			__a	:    Any			   =							(RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
			__a	:    Any			   =							(
			    {"""feature-extraction""": RwkvModel, """text-generation""": RwkvForCausalLM} if is_torch_available() else {}
			)
			# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
			__a	:    Optional[int]			   =							False
			__a	:    Dict			   =							False
			__a	:    Dict			   =							False
			__a	:    Any			   =							False
			__a	:    Dict			   =							False  # Rwkv does not support head masking
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							int       =   RwkvModelTester(self       )
				_UpperCAmelCase      :							List[str]       =   ConfigTester(
				    self ,						config_class=_A ,						n_embd=37 ,						common_properties=["""hidden_size""", """num_hidden_layers"""]       )
			def 							__snake_case	(						self ,						_A ,						_A ,						_A=None       )      ->      Union[str, Any]:






				'''simple docstring'''







				if isinstance(_A ,						torch.Tensor       ):
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Dict       =   member.max().item(), member.min().item()
				elif isinstance(_A ,						_A       ) or isinstance(_A ,						_A       ):
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   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"""       )

				_UpperCAmelCase  ,				_UpperCAmelCase      :							Any       =   container

				_UpperCAmelCase      :							Tuple       =   (min_value >= expected_min) and (max_value <= expected_max)

				if not is_inside_interval:
					_UpperCAmelCase      :							List[Any]       =   """%s not found in %s""" % (safe_repr(_A       ), safe_repr(_A       ))
					self.fail(self._formatMessage(_A ,						_A       )       )
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				self.config_tester.run_common_tests()
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[int]       =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_rwkv_model(*_A       )
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							int       =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_causl_lm(*_A       )
			def 							__snake_case	(						self       )      ->      str:






				'''simple docstring'''







				_UpperCAmelCase      :							Dict       =   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_state_equivalency(*_A       )
			def 							__snake_case	(						self       )      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase      :							str       =   self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
					_UpperCAmelCase      :							List[Any]       =   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 							__snake_case	(						self       )      ->      Any:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase      :							Dict       =   self.model_tester.prepare_config_and_inputs_for_common()
				_UpperCAmelCase      :							List[Any]       =   True

				_UpperCAmelCase      :							Tuple       =   getattr(self.model_tester ,						"""seq_length""" ,						_A       )

				for model_class in self.all_model_classes:
					_UpperCAmelCase      :							int       =   True
					_UpperCAmelCase      :							Optional[int]       =   False
					_UpperCAmelCase      :							List[Any]       =   True
					_UpperCAmelCase      :							Optional[Any]       =   model_class(_A       )
					model.to(_A       )
					model.eval()

					_UpperCAmelCase      :							str       =   self._prepare_for_class(_A ,						_A       )
					_UpperCAmelCase      :							Dict       =   inputs["""input_ids"""].shape[0]
					with torch.no_grad():
						_UpperCAmelCase      :							Dict       =   model(**_A       )
					_UpperCAmelCase      :							int       =   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"]
					_UpperCAmelCase      :							Any       =   True
					_UpperCAmelCase      :							Optional[Any]       =   model_class(_A       )
					model.to(_A       )
					model.eval()

					_UpperCAmelCase      :							Any       =   self._prepare_for_class(_A ,						_A       )
					_UpperCAmelCase      :							Tuple       =   inputs["""input_ids"""].shape[0]
					with torch.no_grad():
						_UpperCAmelCase      :							Tuple       =   model(**_A       )
					_UpperCAmelCase      :							int       =   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] ,						)
					_UpperCAmelCase      :							Dict       =   len(_A       )

					# Check attention is always last and order is fine
					_UpperCAmelCase      :							Optional[Any]       =   True
					_UpperCAmelCase      :							Union[str, Any]       =   True
					_UpperCAmelCase      :							List[str]       =   model_class(_A       )
					model.to(_A       )
					model.eval()

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

					_UpperCAmelCase      :							int       =   1
					self.assertEqual(out_len + added_hidden_states ,						len(_A       )       )

					_UpperCAmelCase      :							Union[str, Any]       =   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 							__snake_case	(						self       )      ->      str:






				'''simple docstring'''







				for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					_UpperCAmelCase      :							str       =   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       _UpperCAmelCase     (						unittest.TestCase):
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				_UpperCAmelCase      :							int       =   """RWKV/rwkv-4-169m-pile"""
				_UpperCAmelCase      :							Optional[int]       =   AutoTokenizer.from_pretrained(self.model_id       )
			def 							__snake_case	(						self       )      ->      Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase      :							Dict       =   """Hello my name is Jasmine and I am a newbie to the"""
				_UpperCAmelCase      :							List[str]       =   RwkvForCausalLM.from_pretrained(self.model_id       ).to(_A       )

				_UpperCAmelCase      :							Optional[int]       =   self.tokenizer("""Hello my name is""" ,						return_tensors="""pt"""       ).input_ids.to(_A       )
				_UpperCAmelCase      :							Tuple       =   model.generate(_A ,						max_new_tokens=10       )
				_UpperCAmelCase      :							Dict       =   self.tokenizer.decode(output[0].tolist()       )

				self.assertEqual(_A ,						_A       )
			def 							__snake_case	(						self       )      ->      Optional[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[Any]       =   """Hello my name is Jasmine and I am a newbie to the"""

				_UpperCAmelCase      :							Union[str, Any]       =   self.tokenizer("""Hello my name is""" ,						return_tensors="""pt"""       ).input_ids.to(_A       )
				_UpperCAmelCase      :							Union[str, Any]       =   RwkvForCausalLM.from_pretrained(self.model_id ,						torch_dtype=torch.bfloataa       ).to(_A       )

				_UpperCAmelCase      :							Union[str, Any]       =   model.generate(_A ,						max_new_tokens=10       )
				_UpperCAmelCase      :							Any       =   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 ...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


lowerCamelCase__			:					Dict							  =					logging.get_logger(__name__)

lowerCamelCase__			:					Optional[Any]							  =					{
    '''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       _UpperCAmelCase     (						__a):
			__a	:    Optional[int]			   =							"""blenderbot"""
			__a	:    Tuple			   =							["""past_key_values"""]
			__a	:    int			   =							{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}
			def __init__(						self ,						_A=80_08 ,						_A=1_28 ,						_A=2 ,						_A=1_02_40 ,						_A=32 ,						_A=24 ,						_A=1_02_40 ,						_A=32 ,						_A=0.0 ,						_A=0.0 ,						_A=True ,						_A=True ,						_A="gelu" ,						_A=25_60 ,						_A=0.1 ,						_A=0.0 ,						_A=0.0 ,						_A=0.02 ,						_A=1 ,						_A=False ,						_A=0 ,						_A=1 ,						_A=2 ,						_A=3 ,						_A=2 ,						**_A ,						)      ->      Union[str, Any]:






				'''simple docstring'''







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

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






class       _UpperCAmelCase     (						__a):
			@property
			def 							__snake_case	(						self       )      ->      Mapping[str, Mapping[int, str]]:






				'''simple docstring'''







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






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Optional[Any]       =   super().outputs
				else:
					_UpperCAmelCase      :							Optional[int]       =   super(_A ,						self       ).outputs
					if self.use_past:
						_UpperCAmelCase  ,				_UpperCAmelCase      :							int       =   self.num_layers
						for i in range(_A       ):
							_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: """past_sequence + sequence"""}
							_UpperCAmelCase      :							Optional[Any]       =   {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]:






				'''simple docstring'''







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

					for _ in range(_A       ):
						common_inputs["past_key_values"].append(
						    (
						        torch.zeros(_A       ),
						        torch.zeros(_A       ),
						        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]:






				'''simple docstring'''







				_UpperCAmelCase      :							Union[str, Any]       =   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
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Any       =   common_inputs["""input_ids"""].shape
					_UpperCAmelCase      :							Union[str, Any]       =   seqlen
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[Any]       =   self.num_layers
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   self.num_attention_heads
					_UpperCAmelCase      :							Optional[Any]       =   (
					    batch,
					    num_encoder_attention_heads,
					    past_key_values_length,
					    self._config.hidden_size // num_encoder_attention_heads,
					)
					_UpperCAmelCase      :							str       =   common_inputs["""attention_mask"""].dtype
					_UpperCAmelCase      :							Tuple       =   torch.cat(
					    [common_inputs["""attention_mask"""], torch.ones(_A ,						_A ,						dtype=_A       )] ,						dim=1       )
					_UpperCAmelCase      :							str       =   [
					    (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]:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   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
				_UpperCAmelCase      :							List[str]       =   tokenizer.num_special_tokens_to_add(_A       )
				_UpperCAmelCase      :							Tuple       =   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
				_UpperCAmelCase      :							str       =   [""" """.join([tokenizer.unk_token]       ) * seq_length] * batch_size
				_UpperCAmelCase      :							str       =   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]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Any       =   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":
					_UpperCAmelCase      :							int       =   self._generate_dummy_inputs_for_causal_lm(
					    _A ,						batch_size=_A ,						seq_length=_A ,						is_pair=_A ,						framework=_A       )
				else:
					_UpperCAmelCase      :							Tuple       =   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       )      ->      Optional[int]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Union[str, Any]       =   super()._flatten_past_key_values_(_A ,						_A ,						_A ,						_A       )
				else:
					_UpperCAmelCase      :							Tuple       =   super(_A ,						self       )._flatten_past_key_values_(
					    _A ,						_A ,						_A ,						_A       )
			def 							__snake_case	(						self ,						_A ,						_A       )      ->      Optional[Any]:






				'''simple docstring'''







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

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

				_UpperCAmelCase      :							List[str]       =   """past_encoder_sequence"""
				_UpperCAmelCase      :							int       =   """past_decoder_sequence""" if direction == """inputs""" else """past_decoder_sequence + sequence"""

				for i in range(_A       ):
					_UpperCAmelCase      :							Union[str, Any]       =   {0: """batch""", 2: decoder_sequence}
					_UpperCAmelCase      :							Optional[int]       =   {0: """batch""", 2: decoder_sequence}
					_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: encoder_sequence}
					_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: encoder_sequence}





