





import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
     import jax
     import jax.numpy as jnp

     from transformers.generation import (
         FlaxForcedBOSTokenLogitsProcessor,
         FlaxForcedEOSTokenLogitsProcessor,
         FlaxLogitsProcessorList,
         FlaxMinLengthLogitsProcessor,
         FlaxTemperatureLogitsWarper,
         FlaxTopKLogitsWarper,
         FlaxTopPLogitsWarper,
     )

@require_flax
class   A							(						unittest.TestCase    ):



      def 						lowerCAmelCase__     (						self:    Dict		,	_lowerCAmelCase:    int		,	_lowerCAmelCase:    int							)    ->     List[str]:
         '''simple docstring'''

         UpperCAmelCase_           =jnp.ones((batch_size, length)							) / length
         return scores



      def 						lowerCAmelCase__     (						self:    List[str]							)    ->     Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =None
         UpperCAmelCase_           =20

         UpperCAmelCase_           =self._get_uniform_logits(batch_size=2		,	length=_lowerCAmelCase							)

         # tweak scores to not be uniform anymore
         UpperCAmelCase_           =scores.at[1, 5].set((1 / length) + 0.1							)  # peak, 1st batch
         UpperCAmelCase_           =scores.at[1, 10].set((1 / length) - 0.4							)  # valley, 1st batch

         # compute softmax
         UpperCAmelCase_           =jax.nn.softmax(_lowerCAmelCase		,	axis=-1							)

         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=0.5							)
         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=1.3							)

         UpperCAmelCase_           =jax.nn.softmax(temp_dist_warper_sharper(_lowerCAmelCase		,	scores.copy()		,	cur_len=_lowerCAmelCase							)		,	axis=-1							)
         UpperCAmelCase_           =jax.nn.softmax(temp_dist_warper_smoother(_lowerCAmelCase		,	scores.copy()		,	cur_len=_lowerCAmelCase							)		,	axis=-1							)

         # uniform distribution stays uniform
         self.assertTrue(jnp.allclose(probs[0, :]		,	warped_prob_sharp[0, :]		,	atol=1e-3							)							)
         self.assertTrue(jnp.allclose(probs[0, :]		,	warped_prob_smooth[0, :]		,	atol=1e-3							)							)

         # sharp peaks get higher, valleys get lower
         self.assertLess(probs[1, :].max()		,	warped_prob_sharp[1, :].max()							)
         self.assertGreater(probs[1, :].min()		,	warped_prob_sharp[1, :].min()							)

         # smooth peaks get lower, valleys get higher
         self.assertGreater(probs[1, :].max()		,	warped_prob_smooth[1, :].max()							)
         self.assertLess(probs[1, :].min()		,	warped_prob_smooth[1, :].min()							)



      def 						lowerCAmelCase__     (						self:    Optional[Any]							)    ->     int:
         '''simple docstring'''

         UpperCAmelCase_           =None
         UpperCAmelCase_           =10
         UpperCAmelCase_           =2

         # create ramp distribution
         UpperCAmelCase_           =np.broadcast_to(np.arange(_lowerCAmelCase							)[None, :]		,	(batch_size, vocab_size)							).copy()
         UpperCAmelCase_           =ramp_logits[1:, : vocab_size // 2] + vocab_size

         UpperCAmelCase_           =FlaxTopKLogitsWarper(3							)

         UpperCAmelCase_           =top_k_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # check that correct tokens are filtered
         self.assertListEqual(jnp.isinf(scores[0]							).tolist()		,	7 * [True] + 3 * [False]							)
         self.assertListEqual(jnp.isinf(scores[1]							).tolist()		,	2 * [True] + 3 * [False] + 5 * [True]							)

         # check special case
         UpperCAmelCase_           =5
         UpperCAmelCase_           =FlaxTopKLogitsWarper(top_k=1		,	filter_value=0.0		,	min_tokens_to_keep=3							)

         UpperCAmelCase_           =np.broadcast_to(np.arange(_lowerCAmelCase							)[None, :]		,	(batch_size, length)							).copy()
         UpperCAmelCase_           =top_k_warp_safety_check(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
         self.assertListEqual((scores == 0.0).sum(axis=-1							).tolist()		,	[2, 2]							)



      def 						lowerCAmelCase__     (						self:    Tuple							)    ->     List[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =None
         UpperCAmelCase_           =10
         UpperCAmelCase_           =2

         # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
         UpperCAmelCase_           =np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]							)							)

         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.8							)
         UpperCAmelCase_           =np.exp(top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)							)

         # dist should be filtered to keep min num values so that sum is >= top_p
         # exp (-inf) => 0
         UpperCAmelCase_           =np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]							)
         self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # check edge cases with negative and extreme logits
         UpperCAmelCase_           =np.broadcast_to(np.arange(_lowerCAmelCase							)[None, :]		,	(batch_size, vocab_size)							).copy() - (
             vocab_size // 2
         )

         # make ramp_logits more extreme
         UpperCAmelCase_           =ramp_logits[1] * 1_00.0

         # make sure at least 2 tokens are kept
         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.9		,	min_tokens_to_keep=2		,	filter_value=0.0							)
         UpperCAmelCase_           =top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
         self.assertListEqual((filtered_dist != 0.0).sum(axis=-1							).tolist()		,	[3, 2]							)



      def 						lowerCAmelCase__     (						self:    Any							)    ->     Any:
         '''simple docstring'''

         UpperCAmelCase_           =20
         UpperCAmelCase_           =4
         UpperCAmelCase_           =0

         UpperCAmelCase_           =FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=_lowerCAmelCase							)

         # check that min length is applied at length 5
         UpperCAmelCase_           =ids_tensor((batch_size, 20)		,	vocab_size=20							)
         UpperCAmelCase_           =5
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =min_dist_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()		,	4 * [-float("inf"							)]							)

         # check that min length is not applied anymore at length 15
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =15
         UpperCAmelCase_           =min_dist_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertFalse(jnp.isinf(_lowerCAmelCase							).any()							)



      def 						lowerCAmelCase__     (						self:    Any							)    ->     Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_           =20
         UpperCAmelCase_           =4
         UpperCAmelCase_           =0

         UpperCAmelCase_           =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase							)

         # check that all scores are -inf except the bos_token_id score
         UpperCAmelCase_           =ids_tensor((batch_size, 1)		,	vocab_size=20							)
         UpperCAmelCase_           =1
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]							).all()							)
         self.assertListEqual(scores[:, bos_token_id].tolist()		,	4 * [0]							)  # score for bos_token_id shold be zero

         # check that bos_token_id is not forced if current length is greater than 1
         UpperCAmelCase_           =3
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertFalse(jnp.isinf(_lowerCAmelCase							).any()							)



      def 						lowerCAmelCase__     (						self:    str							)    ->     Any:
         '''simple docstring'''

         UpperCAmelCase_           =20
         UpperCAmelCase_           =4
         UpperCAmelCase_           =0
         UpperCAmelCase_           =5

         UpperCAmelCase_           =FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCAmelCase		,	eos_token_id=_lowerCAmelCase							)

         # check that all scores are -inf except the eos_token_id when max_length is reached
         UpperCAmelCase_           =ids_tensor((batch_size, 4)		,	vocab_size=20							)
         UpperCAmelCase_           =4
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]							).all()							)
         self.assertListEqual(scores[:, eos_token_id].tolist()		,	4 * [0]							)  # score for eos_token_id should be zero

         # check that eos_token_id is not forced if max_length is not reached
         UpperCAmelCase_           =3
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertFalse(jnp.isinf(_lowerCAmelCase							).any()							)



      def 						lowerCAmelCase__     (						self:    List[str]							)    ->     Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_           =4
         UpperCAmelCase_           =10
         UpperCAmelCase_           =15
         UpperCAmelCase_           =2
         UpperCAmelCase_           =1
         UpperCAmelCase_           =15

         # dummy input_ids and scores
         UpperCAmelCase_           =ids_tensor((batch_size, sequence_length)		,	_lowerCAmelCase							)
         UpperCAmelCase_           =input_ids.copy()

         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =scores.copy()

         # instantiate all dist processors
         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=0.5							)
         UpperCAmelCase_           =FlaxTopKLogitsWarper(3							)
         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.8							)

         # instantiate all logits processors
         UpperCAmelCase_           =FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCAmelCase		,	eos_token_id=_lowerCAmelCase							)

         UpperCAmelCase_           =10

         # no processor list
         UpperCAmelCase_           =temp_dist_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =top_k_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =min_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =bos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =eos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # with processor list
         UpperCAmelCase_           =FlaxLogitsProcessorList(
             [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
         UpperCAmelCase_           =processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # scores should be equal
         self.assertTrue(jnp.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()		,	input_ids_comp.tolist()							)






      def 						lowerCAmelCase__     (						self:    str							)    ->     Tuple:
         '''simple docstring'''

         UpperCAmelCase_           =4
         UpperCAmelCase_           =10
         UpperCAmelCase_           =15
         UpperCAmelCase_           =2
         UpperCAmelCase_           =1
         UpperCAmelCase_           =15

         # dummy input_ids and scores
         UpperCAmelCase_           =ids_tensor((batch_size, sequence_length)		,	_lowerCAmelCase							)
         UpperCAmelCase_           =input_ids.copy()

         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =scores.copy()

         # instantiate all dist processors
         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=0.5							)
         UpperCAmelCase_           =FlaxTopKLogitsWarper(3							)
         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.8							)

         # instantiate all logits processors
         UpperCAmelCase_           =FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCAmelCase		,	eos_token_id=_lowerCAmelCase							)

         UpperCAmelCase_           =10

         # no processor list
         def run_no_processor_list(_lowerCAmelCase:    int		,	_lowerCAmelCase:    Optional[Any]		,	_lowerCAmelCase:    Dict							):
            UpperCAmelCase_           =temp_dist_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =top_k_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =min_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =bos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =eos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            return scores

         # with processor list
         def run_processor_list(_lowerCAmelCase:    Optional[int]		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    Optional[int]							):
            UpperCAmelCase_           =FlaxLogitsProcessorList(
                [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
            UpperCAmelCase_           =processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            return scores

         UpperCAmelCase_           =jax.jit(_lowerCAmelCase							)
         UpperCAmelCase_           =jax.jit(_lowerCAmelCase							)

         UpperCAmelCase_           =jitted_run_no_processor_list(_lowerCAmelCase		,	_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =jitted_run_processor_list(_lowerCAmelCase		,	_lowerCAmelCase		,	_lowerCAmelCase							)

         # scores should be equal
         self.assertTrue(jnp.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()		,	input_ids_comp.tolist()							)







'''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


UpperCAmelCase					:			str       =	logging.get_logger(__name__)

UpperCAmelCase					:			Optional[Any]       =	{
    'facebook/blenderbot_small-90M': 'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json',
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}







class 		lowerCAmelCase__  ( a   ):



 """simple docstring"""



 lowerCAmelCase__       = "blenderbot-small"
 lowerCAmelCase__       = ["past_key_values"]
 lowerCAmelCase__       = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}
 def __init__(						self   :			List[Any]   ,				__SCREAMING_SNAKE_CASE   :			List[Any]=50_265   ,				__SCREAMING_SNAKE_CASE   :			Tuple=512   ,				__SCREAMING_SNAKE_CASE   :			Tuple=8   ,				__SCREAMING_SNAKE_CASE   :			str=2_048   ,				__SCREAMING_SNAKE_CASE   :			Tuple=16   ,				__SCREAMING_SNAKE_CASE   :			Dict=8   ,				__SCREAMING_SNAKE_CASE   :			Dict=2_048   ,				__SCREAMING_SNAKE_CASE   :			str=16   ,				__SCREAMING_SNAKE_CASE   :			Tuple=0.0   ,				__SCREAMING_SNAKE_CASE   :			List[str]=0.0   ,				__SCREAMING_SNAKE_CASE   :			Any=True   ,				__SCREAMING_SNAKE_CASE   :			Any=True   ,				__SCREAMING_SNAKE_CASE   :			List[Any]="gelu"   ,				__SCREAMING_SNAKE_CASE   :			Tuple=512   ,				__SCREAMING_SNAKE_CASE   :			List[str]=0.1   ,				__SCREAMING_SNAKE_CASE   :			List[str]=0.0   ,				__SCREAMING_SNAKE_CASE   :			Optional[int]=0.0   ,				__SCREAMING_SNAKE_CASE   :			List[Any]=0.02   ,				__SCREAMING_SNAKE_CASE   :			Any=1   ,				__SCREAMING_SNAKE_CASE   :			Any=False   ,				__SCREAMING_SNAKE_CASE   :			Optional[int]=0   ,				__SCREAMING_SNAKE_CASE   :			Union[str, Any]=1   ,				__SCREAMING_SNAKE_CASE   :			Optional[int]=2   ,				__SCREAMING_SNAKE_CASE   :			Dict=2   ,				**__SCREAMING_SNAKE_CASE   :			Any   ,				)    ->					Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = vocab_size
      __SCREAMING_SNAKE_CASE      = max_position_embeddings
      __SCREAMING_SNAKE_CASE      = d_model
      __SCREAMING_SNAKE_CASE      = encoder_ffn_dim
      __SCREAMING_SNAKE_CASE      = encoder_layers
      __SCREAMING_SNAKE_CASE      = encoder_attention_heads
      __SCREAMING_SNAKE_CASE      = decoder_ffn_dim
      __SCREAMING_SNAKE_CASE      = decoder_layers
      __SCREAMING_SNAKE_CASE      = decoder_attention_heads
      __SCREAMING_SNAKE_CASE      = dropout
      __SCREAMING_SNAKE_CASE      = attention_dropout
      __SCREAMING_SNAKE_CASE      = activation_dropout
      __SCREAMING_SNAKE_CASE      = activation_function
      __SCREAMING_SNAKE_CASE      = init_std
      __SCREAMING_SNAKE_CASE      = encoder_layerdrop
      __SCREAMING_SNAKE_CASE      = decoder_layerdrop
      __SCREAMING_SNAKE_CASE      = use_cache
      __SCREAMING_SNAKE_CASE      = encoder_layers
      __SCREAMING_SNAKE_CASE      = scale_embedding  # scale factor will be sqrt(d_model) if True

      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   ,				forced_eos_token_id=__SCREAMING_SNAKE_CASE   ,				**__SCREAMING_SNAKE_CASE   ,				)







class 		lowerCAmelCase__  ( a   ):



 """simple docstring"""



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


      """simple docstring"""




      if self.task in ["default", "seq2seq-lm"]:
           __SCREAMING_SNAKE_CASE      = OrderedDict(
               [
                   ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                   ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
               ]     )

           if self.use_past:
                __SCREAMING_SNAKE_CASE      = {0: """batch"""}
                __SCREAMING_SNAKE_CASE      = {0: """batch""", 1: """past_decoder_sequence + sequence"""}
           else:
                __SCREAMING_SNAKE_CASE      = {0: """batch""", 1: """decoder_sequence"""}
                __SCREAMING_SNAKE_CASE      = {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.
           __SCREAMING_SNAKE_CASE      = OrderedDict(
               [
                   ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                   ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
               ]     )
           if self.use_past:
                __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = self.num_layers
                for i in range(__SCREAMING_SNAKE_CASE     ):
                     __SCREAMING_SNAKE_CASE      = {0: """batch""", 2: """past_sequence + sequence"""}
                     __SCREAMING_SNAKE_CASE      = {0: """batch""", 2: """past_sequence + sequence"""}
      else:
           __SCREAMING_SNAKE_CASE      = 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  UpperCAmelCase__						(						self   :			Dict     )    ->					Mapping[str, Mapping[int, str]]:


      """simple docstring"""




      if self.task in ["default", "seq2seq-lm"]:
           __SCREAMING_SNAKE_CASE      = super().outputs
      else:
           __SCREAMING_SNAKE_CASE      = super(__SCREAMING_SNAKE_CASE   ,				self     ).outputs
           if self.use_past:
                __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = self.num_layers
                for i in range(__SCREAMING_SNAKE_CASE     ):
                     __SCREAMING_SNAKE_CASE      = {0: """batch""", 2: """past_sequence + sequence"""}
                     __SCREAMING_SNAKE_CASE      = {0: """batch""", 2: """past_sequence + sequence"""}
      return common_outputs
 def  UpperCAmelCase__						(						self   :			int   ,				__SCREAMING_SNAKE_CASE   :			PreTrainedTokenizer   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			bool = False   ,				__SCREAMING_SNAKE_CASE   :			Optional[TensorType] = None   ,				)    ->					Mapping[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = 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
      __SCREAMING_SNAKE_CASE      = seq_length if not self.use_past else 1
      __SCREAMING_SNAKE_CASE      = 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     )
      __SCREAMING_SNAKE_CASE      = {f'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
      __SCREAMING_SNAKE_CASE      = 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
           __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = common_inputs["""input_ids"""].shape
           __SCREAMING_SNAKE_CASE      = common_inputs["""decoder_input_ids"""].shape[1]
           __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = self.num_attention_heads
           __SCREAMING_SNAKE_CASE      = (
               batch,
               num_encoder_attention_heads,
               encoder_seq_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )
           __SCREAMING_SNAKE_CASE      = decoder_seq_length + 3
           __SCREAMING_SNAKE_CASE      = (
               batch,
               num_decoder_attention_heads,
               decoder_past_length,
               self._config.hidden_size // num_decoder_attention_heads,
           )

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

           __SCREAMING_SNAKE_CASE      = []
           # If the number of encoder and decoder layers are present in the model configuration, both are considered
           __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = self.num_layers
           __SCREAMING_SNAKE_CASE      = min(__SCREAMING_SNAKE_CASE   ,				__SCREAMING_SNAKE_CASE     )
           __SCREAMING_SNAKE_CASE      = max(__SCREAMING_SNAKE_CASE   ,				__SCREAMING_SNAKE_CASE     ) - min_num_layers
           __SCREAMING_SNAKE_CASE      = """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.
           __SCREAMING_SNAKE_CASE      = 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  UpperCAmelCase__						(						self   :			Dict   ,				__SCREAMING_SNAKE_CASE   :			PreTrainedTokenizer   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			bool = False   ,				__SCREAMING_SNAKE_CASE   :			Optional[TensorType] = None   ,				)    ->					Mapping[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = 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
           __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = common_inputs["""input_ids"""].shape
           # Not using the same length for past_key_values
           __SCREAMING_SNAKE_CASE      = seqlen + 2
           __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = self.num_layers
           __SCREAMING_SNAKE_CASE  ,     __SCREAMING_SNAKE_CASE      = self.num_attention_heads
           __SCREAMING_SNAKE_CASE      = (
               batch,
               num_encoder_attention_heads,
               past_key_values_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )

           __SCREAMING_SNAKE_CASE      = common_inputs["""attention_mask"""].dtype
           __SCREAMING_SNAKE_CASE      = torch.cat(
               [common_inputs["""attention_mask"""], torch.ones(__SCREAMING_SNAKE_CASE   ,				__SCREAMING_SNAKE_CASE   ,				dtype=__SCREAMING_SNAKE_CASE     )]   ,				dim=1     )
           __SCREAMING_SNAKE_CASE      = [
               (torch.zeros(__SCREAMING_SNAKE_CASE     ), torch.zeros(__SCREAMING_SNAKE_CASE     )) for _ in range(__SCREAMING_SNAKE_CASE     )
           ]
      return common_inputs
 def  UpperCAmelCase__						(						self   :			Optional[Any]   ,				__SCREAMING_SNAKE_CASE   :			PreTrainedTokenizer   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			bool = False   ,				__SCREAMING_SNAKE_CASE   :			Optional[TensorType] = None   ,				)    ->					Mapping[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = 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
      __SCREAMING_SNAKE_CASE      = tokenizer.num_special_tokens_to_add(__SCREAMING_SNAKE_CASE     )
      __SCREAMING_SNAKE_CASE      = 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
      __SCREAMING_SNAKE_CASE      = [""" """.join([tokenizer.unk_token]     ) * seq_length] * batch_size
      __SCREAMING_SNAKE_CASE      = dict(tokenizer(__SCREAMING_SNAKE_CASE   ,				return_tensors=__SCREAMING_SNAKE_CASE     )     )
      return common_inputs
 def  UpperCAmelCase__						(						self   :			Any   ,				__SCREAMING_SNAKE_CASE   :			PreTrainedTokenizer   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			int = -1   ,				__SCREAMING_SNAKE_CASE   :			bool = False   ,				__SCREAMING_SNAKE_CASE   :			Optional[TensorType] = None   ,				)    ->					Mapping[str, Any]:


      """simple docstring"""




      if self.task in ["default", "seq2seq-lm"]:
           __SCREAMING_SNAKE_CASE      = 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":
           __SCREAMING_SNAKE_CASE      = 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:
           __SCREAMING_SNAKE_CASE      = 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  UpperCAmelCase__						(						self   :			Dict   ,				__SCREAMING_SNAKE_CASE   :			str   ,				__SCREAMING_SNAKE_CASE   :			Dict   ,				__SCREAMING_SNAKE_CASE   :			List[str]   ,				__SCREAMING_SNAKE_CASE   :			Union[str, Any]     )    ->					Optional[Any]:


      """simple docstring"""




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






