





'''simple docstring'''


import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch


def 						__a(SCREAMING_SNAKE_CASE_	:       Optional[int]     ):  # picklable for multiprocessing



    '''simple docstring'''





    return x.sum()


def 						__a(SCREAMING_SNAKE_CASE_	:       List[Any]     ):  # picklable for multiprocessing



    '''simple docstring'''





    return i + 1

@dataclass
class  lowerCAmelCase_		:
 __lowerCamelCase    :				int
 __lowerCamelCase    :				str







class  lowerCAmelCase_		(		__magic_name__		):




 def 				_snake_case (						self		)     ->				List[str]:
     _lowerCAmelCase						      =   {}
     _lowerCAmelCase						      =   []
     _lowerCAmelCase						      =   1
     _lowerCAmelCase						      =   [1, 2]
     _lowerCAmelCase						      =   {"a": 1, "b": 2}
     _lowerCAmelCase						      =   {"a": [1, 2], "b": [3, 4]}
     _lowerCAmelCase						      =   {"a": {"1": 1}, "b": 2}
     _lowerCAmelCase						      =   {"a": 1, "b": 2, "c": 3, "d": 4}
     _lowerCAmelCase						      =   {}
     _lowerCAmelCase						      =   []
     _lowerCAmelCase						      =   2
     _lowerCAmelCase						      =   [2, 3]
     _lowerCAmelCase						      =   {"a": 2, "b": 3}
     _lowerCAmelCase						      =   {"a": [2, 3], "b": [4, 5]}
     _lowerCAmelCase						      =   {"a": {"1": 2}, "b": 3}
     _lowerCAmelCase						      =   {"a": 2, "b": 3, "c": 4, "d": 5}
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase		)      ,  _lowerCAmelCase		)

     _lowerCAmelCase						      =   2
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)

     _lowerCAmelCase						      =   {"a": np.eye(2		), "b": np.zeros(3		), "c": np.ones(2		)}
     _lowerCAmelCase						      =   {"a": 2, "b": 0, "c": 2}
     _lowerCAmelCase						      =   {
         "a": np.eye(2		).astype(_lowerCAmelCase		),
         "b": np.zeros(3		).astype(_lowerCAmelCase		),
         "c": np.ones(2		).astype(_lowerCAmelCase		),
     }
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  map_numpy=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(
         {k: v.tolist() for k, v in map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  map_numpy=_lowerCAmelCase		).items()}      ,  {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,  )
     self.assertEqual(map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  map_numpy=_lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)      ,  _lowerCAmelCase		)
     self.assertEqual(
         {k: v.tolist() for k, v in map_nested(_lowerCAmelCase      ,  _lowerCAmelCase      ,  map_numpy=_lowerCAmelCase      ,  num_proc=_lowerCAmelCase		).items()}      ,  {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,  )
     with self.assertRaises(_lowerCAmelCase		):  # can't pickle a local lambda
         map_nested(lambda _lowerCAmelCase		: x + 1      ,  _lowerCAmelCase      ,  num_proc=_lowerCAmelCase		)




 def 				_snake_case (						self		)     ->				Union[str, Any]:
     _lowerCAmelCase						      =   {"a": 1, "b": 2}
     _lowerCAmelCase						      =   {"a": 3, "b": 4}
     _lowerCAmelCase						      =   {"a": 5, "b": 6}
     _lowerCAmelCase						      =   sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]		)
     self.assertEqual(sorted(zip_dict(_lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)		)      ,  _lowerCAmelCase		)





 def 				_snake_case (						self		)     ->				str:

     class  lowerCAmelCase_		:
      __lowerCamelCase    :				Tuple								=       "bar"







     _lowerCAmelCase						      =   Foo()
     self.assertEqual(foo.my_attr      ,  "bar"		)
     with temporary_assignment(_lowerCAmelCase      ,  "my_attr"      ,  "BAR"		):
         self.assertEqual(foo.my_attr      ,  "BAR"		)
     self.assertEqual(foo.my_attr      ,  "bar"		)







@pytest.mark.parametrize(
    "iterable_length, num_proc, expected_num_proc"		,							[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]		,							)
def 						__a(SCREAMING_SNAKE_CASE_	:       int		,							SCREAMING_SNAKE_CASE_	:       int		,							SCREAMING_SNAKE_CASE_	:       Any     ):



    '''simple docstring'''





    with patch("datasets.utils.py_utils._single_map_nested"     ) as mock_single_map_nested, patch(
        "datasets.parallel.parallel.Pool"     ) as mock_multiprocessing_pool:
        _lowerCAmelCase						      =   {F'''{i}''': i for i in range(SCREAMING_SNAKE_CASE_     )}
        _lowerCAmelCase						      =   map_nested(lambda SCREAMING_SNAKE_CASE_     : x + 10		,							SCREAMING_SNAKE_CASE_		,							num_proc=SCREAMING_SNAKE_CASE_		,							parallel_min_length=16     )
        if expected_num_proc == 1:
            assert mock_single_map_nested.called
            assert not mock_multiprocessing_pool.called
        else:
            assert not mock_single_map_nested.called
            assert mock_multiprocessing_pool.called
            assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc

class  lowerCAmelCase_		(		__magic_name__		):




 @require_tf
 def 				_snake_case (						self		)     ->				Optional[Any]:
     import tensorflow as tf
     from tensorflow.keras import layers

     _lowerCAmelCase						      =   layers.Dense(2		)

     def gen_random_output():
         _lowerCAmelCase						      =   tf.random.uniform((1, 3)		)
         return model(_lowerCAmelCase		).numpy()

     with temp_seed(42      ,  set_tensorflow=_lowerCAmelCase		):
         _lowerCAmelCase						      =   gen_random_output()
     with temp_seed(42      ,  set_tensorflow=_lowerCAmelCase		):
         _lowerCAmelCase						      =   gen_random_output()
     _lowerCAmelCase						      =   gen_random_output()

     np.testing.assert_equal(_lowerCAmelCase      ,  _lowerCAmelCase		)
     self.assertGreater(np.abs(outa - outa		).sum()      ,  0		)




 @require_torch
 def 				_snake_case (						self		)     ->				Any:
     import torch

     def gen_random_output():
         _lowerCAmelCase						      =   torch.nn.Linear(3      ,  2		)
         _lowerCAmelCase						      =   torch.rand(1      ,  3		)
         return model(_lowerCAmelCase		).detach().numpy()

     with temp_seed(42      ,  set_pytorch=_lowerCAmelCase		):
         _lowerCAmelCase						      =   gen_random_output()
     with temp_seed(42      ,  set_pytorch=_lowerCAmelCase		):
         _lowerCAmelCase						      =   gen_random_output()
     _lowerCAmelCase						      =   gen_random_output()

     np.testing.assert_equal(_lowerCAmelCase      ,  _lowerCAmelCase		)
     self.assertGreater(np.abs(outa - outa		).sum()      ,  0		)





 def 				_snake_case (						self		)     ->				Dict:
     def gen_random_output():
         return np.random.rand(1      ,  3		)

     with temp_seed(42		):
         _lowerCAmelCase						      =   gen_random_output()
     with temp_seed(42		):
         _lowerCAmelCase						      =   gen_random_output()
     _lowerCAmelCase						      =   gen_random_output()

     np.testing.assert_equal(_lowerCAmelCase      ,  _lowerCAmelCase		)
     self.assertGreater(np.abs(outa - outa		).sum()      ,  0		)







@pytest.mark.parametrize("input_data"		,							[{}]     )
def 						__a(SCREAMING_SNAKE_CASE_	:       int     ):



    '''simple docstring'''





    _lowerCAmelCase						      =   NestedDataStructure(SCREAMING_SNAKE_CASE_     ).data
    assert output_data == input_data


@pytest.mark.parametrize(
    "data, expected_output"		,							[
        ({}, []),
        ([], []),
        ("foo", ["foo"]),
        (["foo", "bar"], ["foo", "bar"]),
        ([["foo", "bar"]], ["foo", "bar"]),
        ([[["foo"], ["bar"]]], ["foo", "bar"]),
        ([[["foo"], "bar"]], ["foo", "bar"]),
        ({"a": 1, "b": 2}, [1, 2]),
        ({"a": [1, 2], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[1, 2]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [[3, 4]]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, 4]}, [1, 2, 3, 4]),
        ({"a": [[[1], [2]]], "b": [3, [4]]}, [1, 2, 3, 4]),
        ({"a": {"1": 1}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": 2}, [1, 2]),
        ({"a": {"1": [1]}, "b": [2]}, [1, 2]),
    ]		,							)
def 						__a(SCREAMING_SNAKE_CASE_	:       Tuple		,							SCREAMING_SNAKE_CASE_	:       Union[str, Any]     ):



    '''simple docstring'''





    _lowerCAmelCase						      =   NestedDataStructure(SCREAMING_SNAKE_CASE_     ).flatten()
    assert output == expected_output


def 						__a():



    '''simple docstring'''





    _lowerCAmelCase						      =   A(x=1		,							y="foobar"     )
    _lowerCAmelCase						      =   {"x": 1, "y": "foobar"}
    assert asdict(SCREAMING_SNAKE_CASE_     ) == expected_output

    _lowerCAmelCase						      =   {"a": {"b": A(x=10		,							y="foo"     )}, "c": [A(x=20		,							y="bar"     )]}
    _lowerCAmelCase						      =   {"a": {"b": {"x": 10, "y": "foo"}}, "c": [{"x": 20, "y": "bar"}]}
    assert asdict(SCREAMING_SNAKE_CASE_     ) == expected_output

    with pytest.raises(SCREAMING_SNAKE_CASE_     ):
        asdict([1, A(x=10		,							y="foo"     )]     )


def 						__a(SCREAMING_SNAKE_CASE_	:       str     ):



    '''simple docstring'''





    return text.split()


def 						__a(SCREAMING_SNAKE_CASE_	:       Dict     ):



    '''simple docstring'''





    yield (time.time(), content)
    time.sleep(2     )
    yield (time.time(), content)





def 						__a():



    '''simple docstring'''





    with Pool(2     ) as pool:
        _lowerCAmelCase						      =   list(iflatmap_unordered(SCREAMING_SNAKE_CASE_		,							_split_text		,							kwargs_iterable=[{"text": "hello there"}] * 10     )     )
        assert out.count("hello"     ) == 10
        assert out.count("there"     ) == 10
        assert len(SCREAMING_SNAKE_CASE_     ) == 20

    # check multiprocess from pathos (uses dill for pickling)
    with multiprocess.Pool(2     ) as pool:
        _lowerCAmelCase						      =   list(iflatmap_unordered(SCREAMING_SNAKE_CASE_		,							_split_text		,							kwargs_iterable=[{"text": "hello there"}] * 10     )     )
        assert out.count("hello"     ) == 10
        assert out.count("there"     ) == 10
        assert len(SCREAMING_SNAKE_CASE_     ) == 20

    # check that we get items as fast as possible
    with Pool(2     ) as pool:
        _lowerCAmelCase						      =   []
        for yield_time, content in iflatmap_unordered(
            SCREAMING_SNAKE_CASE_		,							_aseconds_generator_of_aitems_with_timing		,							kwargs_iterable=[{"content": "a"}, {"content": "b"}]     ):
            assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
            out.append(SCREAMING_SNAKE_CASE_     )
        assert out.count("a"     ) == 2
        assert out.count("b"     ) == 2
        assert len(SCREAMING_SNAKE_CASE_     ) == 4









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


_SCREAMING_SNAKE_CASE									=				logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE									=				{
    "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_		(		__magic_name__		):
 __lowerCamelCase    :				Union[str, Any]								=       "blenderbot-small"
 __lowerCamelCase    :				Optional[Any]								=       ["past_key_values"]
 __lowerCamelCase    :				str								=       {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




 def __init__(						self      ,  _lowerCAmelCase=50265      ,  _lowerCAmelCase=512      ,  _lowerCAmelCase=8      ,  _lowerCAmelCase=2048      ,  _lowerCAmelCase=16      ,  _lowerCAmelCase=8      ,  _lowerCAmelCase=2048      ,  _lowerCAmelCase=16      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase="gelu"      ,  _lowerCAmelCase=512      ,  _lowerCAmelCase=0.1      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=0.02      ,  _lowerCAmelCase=1      ,  _lowerCAmelCase=False      ,  _lowerCAmelCase=0      ,  _lowerCAmelCase=1      ,  _lowerCAmelCase=2      ,  _lowerCAmelCase=2      ,  **_lowerCAmelCase      ,  )     ->				Dict:
     _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						      =   use_cache
     _lowerCAmelCase						      =   encoder_layers
     _lowerCAmelCase						      =   scale_embedding  # scale factor will be sqrt(d_model) if True

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







class  lowerCAmelCase_		(		__magic_name__		):




 @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_(_lowerCAmelCase      ,  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(_lowerCAmelCase		):
                 _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(_lowerCAmelCase      ,  self		).outputs
         if self.use_past:
             _lowerCAmelCase      ,      _lowerCAmelCase						      =   self.num_layers
             for i in range(_lowerCAmelCase		):
                 _lowerCAmelCase						      =   {0: "batch", 2: "past_sequence + sequence"}
                 _lowerCAmelCase						      =   {0: "batch", 2: "past_sequence + sequence"}
     return common_outputs




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = False      ,  _lowerCAmelCase = None      ,  )     ->				Mapping[str, Any]:
     _lowerCAmelCase						      =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)

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

     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(_lowerCAmelCase      ,  _lowerCAmelCase		)]      ,  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(_lowerCAmelCase      ,  _lowerCAmelCase		)
         _lowerCAmelCase						      =   max(_lowerCAmelCase      ,  _lowerCAmelCase		) - min_num_layers
         _lowerCAmelCase						      =   "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

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




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = False      ,  _lowerCAmelCase = None      ,  )     ->				Mapping[str, Any]:
     _lowerCAmelCase						      =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)

     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(_lowerCAmelCase      ,  _lowerCAmelCase      ,  dtype=_lowerCAmelCase		)]      ,  dim=1		)
         _lowerCAmelCase						      =   [
             (torch.zeros(_lowerCAmelCase		), torch.zeros(_lowerCAmelCase		)) for _ in range(_lowerCAmelCase		)
         ]
     return common_inputs




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = False      ,  _lowerCAmelCase = 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(
         _lowerCAmelCase      ,  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(_lowerCAmelCase		)
     _lowerCAmelCase						      =   compute_effective_axis_dimension(
         _lowerCAmelCase      ,  fixed_dimension=OnnxConfig.default_fixed_sequence      ,  num_token_to_add=_lowerCAmelCase		)

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




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

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

     return common_inputs





 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				Dict:
     if self.task in ["default", "seq2seq-lm"]:
         _lowerCAmelCase						      =   super()._flatten_past_key_values_(_lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)
     else:
         _lowerCAmelCase						      =   super(_lowerCAmelCase      ,  self		)._flatten_past_key_values_(
             _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)


