




from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


_lowerCAmelCase    : Optional[Any]		= logging.get_logger(__name__)


if is_vision_available():
  import PIL
class       _UpperCamelCase   (					__snake_case		):
      UpperCAmelCase_          =		["""pixel_values"""]
      def __init__(      self :Dict			,   lowerCamelCase :List[str] = True			,   lowerCamelCase :Optional[int] = None			,   lowerCamelCase :Any = PILImageResampling.BICUBIC			,   lowerCamelCase :Optional[int] = True			,   lowerCamelCase :Union[str, Any] = None			,   lowerCamelCase :Optional[int] = True			,   lowerCamelCase :int = 1 / 255			,   lowerCamelCase :Optional[Any] = True			,   lowerCamelCase :int = None			,   lowerCamelCase :Any = None			,   lowerCamelCase :Optional[Any] = True			,   **lowerCamelCase :Optional[int]			,   )  ->      None:
       super().__init__(**__lowerCamelCase   )
       UpperCAmelCase__							=    size if size is not None else {"""shortest_edge""": 224}
       UpperCAmelCase__							=    get_size_dict(__lowerCamelCase			,   default_to_square=__lowerCamelCase   )
       UpperCAmelCase__							=    crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
       UpperCAmelCase__							=    get_size_dict(__lowerCamelCase   )

       UpperCAmelCase__							=    do_resize
       UpperCAmelCase__							=    size
       UpperCAmelCase__							=    resample
       UpperCAmelCase__							=    do_center_crop
       UpperCAmelCase__							=    crop_size
       UpperCAmelCase__							=    do_rescale
       UpperCAmelCase__							=    rescale_factor
       UpperCAmelCase__							=    do_normalize
       UpperCAmelCase__							=    image_mean if image_mean is not None else OPENAI_CLIP_MEAN
       UpperCAmelCase__							=    image_std if image_std is not None else OPENAI_CLIP_STD
       UpperCAmelCase__							=    do_convert_rgb
      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :str			,   lowerCamelCase :Union[str, Any] = PILImageResampling.BICUBIC			,   lowerCamelCase :Any = None			,   **lowerCamelCase :Any			,   )  ->      np.ndarray:
       UpperCAmelCase__							=    get_size_dict(__lowerCamelCase			,   default_to_square=__lowerCamelCase   )
       UpperCAmelCase__							=    get_resize_output_image_size(
           __lowerCamelCase			,   size=(size["height"], size["width"])			,   default_to_square=__lowerCamelCase   )
       return resize(__lowerCamelCase			,   size=__lowerCamelCase			,   resample=__lowerCamelCase			,   data_format=__lowerCamelCase			,   **__lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :int			,   lowerCamelCase :Optional[int] = None			,   **lowerCamelCase :List[Any]			,   )  ->      np.ndarray:
       UpperCAmelCase__							=    get_size_dict(__lowerCamelCase   )
       return center_crop(__lowerCamelCase			,   size=(size["height"], size["width"])			,   data_format=__lowerCamelCase			,   **__lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Any			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :int = None			,   **lowerCamelCase :Dict			,   )  ->      str:
       return rescale(__lowerCamelCase			,   scale=__lowerCamelCase			,   data_format=__lowerCamelCase			,   **__lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Optional[Any] = None			,   **lowerCamelCase :Any			,   )  ->      np.ndarray:
       return normalize(__lowerCamelCase			,   mean=__lowerCamelCase			,   std=__lowerCamelCase			,   data_format=__lowerCamelCase			,   **__lowerCamelCase   )







      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :str = None			,   lowerCamelCase :List[str] = None			,   lowerCamelCase :List[str] = None			,   lowerCamelCase :Optional[Any] = None			,   lowerCamelCase :List[str] = None			,   lowerCamelCase :Union[str, Any] = None			,   lowerCamelCase :Optional[int] = None			,   lowerCamelCase :Union[str, Any] = None			,   lowerCamelCase :Optional[Any] = None			,   lowerCamelCase :Union[str, Any] = None			,   lowerCamelCase :Dict = None			,   lowerCamelCase :Union[str, Any] = None			,   lowerCamelCase :List[Any] = ChannelDimension.FIRST			,   **lowerCamelCase :Optional[Any]			,   )  ->      PIL.Image.Image:
       UpperCAmelCase__							=    do_resize if do_resize is not None else self.do_resize
       UpperCAmelCase__							=    size if size is not None else self.size
       UpperCAmelCase__							=    get_size_dict(__lowerCamelCase			,   default_to_square=__lowerCamelCase   )
       UpperCAmelCase__							=    resample if resample is not None else self.resample
       UpperCAmelCase__							=    do_center_crop if do_center_crop is not None else self.do_center_crop
       UpperCAmelCase__							=    crop_size if crop_size is not None else self.crop_size
       UpperCAmelCase__							=    get_size_dict(__lowerCamelCase   )
       UpperCAmelCase__							=    do_rescale if do_rescale is not None else self.do_rescale
       UpperCAmelCase__							=    rescale_factor if rescale_factor is not None else self.rescale_factor
       UpperCAmelCase__							=    do_normalize if do_normalize is not None else self.do_normalize
       UpperCAmelCase__							=    image_mean if image_mean is not None else self.image_mean
       UpperCAmelCase__							=    image_std if image_std is not None else self.image_std
       UpperCAmelCase__							=    do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

       UpperCAmelCase__							=    make_list_of_images(__lowerCamelCase   )

       if not valid_images(__lowerCamelCase   ):
        raise ValueError(
            "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
            "torch.Tensor, tf.Tensor or jax.ndarray."   )

       if do_resize and size is None:
        raise ValueError("Size must be specified if do_resize is True."   )

       if do_center_crop and crop_size is None:
        raise ValueError("Crop size must be specified if do_center_crop is True."   )

       if do_rescale and rescale_factor is None:
        raise ValueError("Rescale factor must be specified if do_rescale is True."   )

       if do_normalize and (image_mean is None or image_std is None):
        raise ValueError("Image mean and std must be specified if do_normalize is True."   )

       # PIL RGBA images are converted to RGB
       if do_convert_rgb:
        UpperCAmelCase__							=    [convert_to_rgb(__lowerCamelCase   ) for image in images]

       # All transformations expect numpy arrays.
       UpperCAmelCase__							=    [to_numpy_array(__lowerCamelCase   ) for image in images]

       if do_resize:
        UpperCAmelCase__							=    [self.resize(image=__lowerCamelCase			,   size=__lowerCamelCase			,   resample=__lowerCamelCase   ) for image in images]

       if do_center_crop:
        UpperCAmelCase__							=    [self.center_crop(image=__lowerCamelCase			,   size=__lowerCamelCase   ) for image in images]

       if do_rescale:
        UpperCAmelCase__							=    [self.rescale(image=__lowerCamelCase			,   scale=__lowerCamelCase   ) for image in images]

       if do_normalize:
        UpperCAmelCase__							=    [self.normalize(image=__lowerCamelCase			,   mean=__lowerCamelCase			,   std=__lowerCamelCase   ) for image in images]

       UpperCAmelCase__							=    [to_channel_dimension_format(__lowerCamelCase			,   __lowerCamelCase   ) for image in images]

       UpperCAmelCase__							=    {"""pixel_values""": images}
       return BatchFeature(data=__lowerCamelCase			,   tensor_type=__lowerCamelCase   )






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

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


__lowerCAmelCase    :  Optional[int]										= logging.get_logger(__name__)

__lowerCAmelCase    :  Dict										= {
    'google/bigbird-pegasus-large-arxiv': (
        'https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json'
    ),
    'google/bigbird-pegasus-large-pubmed': (
        'https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json'
    ),
    'google/bigbird-pegasus-large-bigpatent': (
        'https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json'
    ),
    # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
}
class 				lowerCamelCase			(							__snake_case					):
      __lowerCamelCase							=					'bigbird_pegasus'
      __lowerCamelCase							=					['past_key_values']
      __lowerCamelCase							=					{
          'num_attention_heads': 'encoder_attention_heads',
          'hidden_size': 'd_model',
          'attention_probs_dropout_prob': 'attention_dropout',
      }






      def __init__(      self  ,		__lowerCamelCase=9_61_03  ,		__lowerCamelCase=40_96  ,		__lowerCamelCase=16  ,		__lowerCamelCase=40_96  ,		__lowerCamelCase=16  ,		__lowerCamelCase=16  ,		__lowerCamelCase=40_96  ,		__lowerCamelCase=16  ,		__lowerCamelCase=0.0  ,		__lowerCamelCase=0.0  ,		__lowerCamelCase=True  ,		__lowerCamelCase=True  ,		__lowerCamelCase="gelu_new"  ,		__lowerCamelCase=10_24  ,		__lowerCamelCase=0.1  ,		__lowerCamelCase=0.0  ,		__lowerCamelCase=0.0  ,		__lowerCamelCase=0.02  ,		__lowerCamelCase=2  ,		__lowerCamelCase=0.0  ,		__lowerCamelCase=True  ,		__lowerCamelCase=0  ,		__lowerCamelCase=2  ,		__lowerCamelCase=1  ,		__lowerCamelCase="block_sparse"  ,		__lowerCamelCase=64  ,		__lowerCamelCase=3  ,		__lowerCamelCase=False  ,		**__lowerCamelCase  ,		)     ->				Dict:


          '''simple docstring'''







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

          # extra config
          snake_case:						Union[str, Any]         =   attention_type
          snake_case:						Optional[int]         =   block_size
          snake_case:						Optional[int]         =   num_random_blocks
          snake_case:						List[Any]         =   use_bias

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





class 				lowerCamelCase			(							__snake_case					):






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


          '''simple docstring'''







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

              if self.use_past:
                  snake_case:						List[Any]         =   {0: """batch"""}
                  snake_case:						Optional[int]         =   {0: """batch""", 1: """past_decoder_sequence + sequence"""}
              else:
                  snake_case:						List[Any]         =   {0: """batch""", 1: """decoder_sequence"""}
                  snake_case:						Union[str, Any]         =   {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.
              snake_case:						List[str]         =   OrderedDict(
                  [
                      ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                      ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                  ]						)
              if self.use_past:
                  snake_case    , snake_case:						str         =   self.num_layers
                  for i in range(__lowerCamelCase						):
                      snake_case:						Any         =   {0: """batch""", 2: """past_sequence + sequence"""}
                      snake_case:						str         =   {0: """batch""", 2: """past_sequence + sequence"""}
          else:
              snake_case:						str         =   OrderedDict(
                  [
                      ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                      ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                      ("""decoder_input_ids""", {0: """batch""", 1: """decoder_sequence"""}),
                      ("""decoder_attention_mask""", {0: """batch""", 1: """decoder_sequence"""}),
                  ]						)

          return common_inputs






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


          '''simple docstring'''







          if self.task in ["default", "seq2seq-lm"]:
              snake_case:						List[str]         =   super().outputs
          else:
              snake_case:						Optional[Any]         =   super(__lowerCamelCase  ,		self						).outputs
              if self.use_past:
                  snake_case    , snake_case:						str         =   self.num_layers
                  for i in range(__lowerCamelCase						):
                      snake_case:						List[Any]         =   {0: """batch""", 2: """past_sequence + sequence"""}
                      snake_case:						int         =   {0: """batch""", 2: """past_sequence + sequence"""}
          return common_outputs






      def 			lowerCAmelCase_   (      self  ,		__lowerCamelCase  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = False  ,		__lowerCamelCase = None  ,		)     ->				Mapping[str, Any]:


          '''simple docstring'''







          snake_case:						List[Any]         =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              __lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)

          # Generate decoder inputs
          snake_case:						Any         =   seq_length if not self.use_past else 1
          snake_case:						List[str]         =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              __lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)
          snake_case:						Dict         =   {F"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
          snake_case:						List[str]         =   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
              snake_case    , snake_case:						Optional[int]         =   common_inputs["""input_ids"""].shape
              snake_case:						Optional[Any]         =   common_inputs["""decoder_input_ids"""].shape[1]
              snake_case    , snake_case:						Any         =   self.num_attention_heads
              snake_case:						Optional[int]         =   (
                  batch,
                  num_encoder_attention_heads,
                  encoder_seq_length,
                  self._config.hidden_size // num_encoder_attention_heads,
              )
              snake_case:						List[Any]         =   decoder_seq_length + 3
              snake_case:						Optional[Any]         =   (
                  batch,
                  num_decoder_attention_heads,
                  decoder_past_length,
                  self._config.hidden_size // num_decoder_attention_heads,
              )

              snake_case:						Tuple         =   torch.cat(
                  [common_inputs["""decoder_attention_mask"""], torch.ones(__lowerCamelCase  ,		__lowerCamelCase						)]  ,		dim=1						)

              snake_case:						Tuple         =   []
              # If the number of encoder and decoder layers are present in the model configuration, both are considered
              snake_case    , snake_case:						Union[str, Any]         =   self.num_layers
              snake_case:						List[str]         =   min(__lowerCamelCase  ,		__lowerCamelCase						)
              snake_case:						List[str]         =   max(__lowerCamelCase  ,		__lowerCamelCase						) - min_num_layers
              snake_case:						Union[str, Any]         =   """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.
              snake_case:						Dict         =   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 			lowerCAmelCase_   (      self  ,		__lowerCamelCase  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = False  ,		__lowerCamelCase = None  ,		)     ->				Mapping[str, Any]:


          '''simple docstring'''







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

              snake_case:						int         =   common_inputs["""attention_mask"""].dtype
              snake_case:						Union[str, Any]         =   torch.cat(
                  [common_inputs["""attention_mask"""], torch.ones(__lowerCamelCase  ,		__lowerCamelCase  ,		dtype=__lowerCamelCase						)]  ,		dim=1						)
              snake_case:						Dict         =   [
                  (torch.zeros(__lowerCamelCase						), torch.zeros(__lowerCamelCase						)) for _ in range(__lowerCamelCase						)
              ]
          return common_inputs






      def 			lowerCAmelCase_   (      self  ,		__lowerCamelCase  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = False  ,		__lowerCamelCase = None  ,		)     ->				Mapping[str, Any]:


          '''simple docstring'''







          snake_case:						str         =   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
          snake_case:						str         =   tokenizer.num_special_tokens_to_add(__lowerCamelCase						)
          snake_case:						int         =   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
          snake_case:						Dict         =   [""" """.join([tokenizer.unk_token]						) * seq_length] * batch_size
          snake_case:						Any         =   dict(tokenizer(__lowerCamelCase  ,		return_tensors=__lowerCamelCase						)						)
          return common_inputs






      def 			lowerCAmelCase_   (      self  ,		__lowerCamelCase  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = -1  ,		__lowerCamelCase = False  ,		__lowerCamelCase = None  ,		)     ->				Mapping[str, Any]:


          '''simple docstring'''







          if self.task in ["default", "seq2seq-lm"]:
              snake_case:						Union[str, Any]         =   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":
              snake_case:						Union[str, Any]         =   self._generate_dummy_inputs_for_causal_lm(
                  __lowerCamelCase  ,		batch_size=__lowerCamelCase  ,		seq_length=__lowerCamelCase  ,		is_pair=__lowerCamelCase  ,		framework=__lowerCamelCase						)
          else:
              snake_case:						int         =   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 			lowerCAmelCase_   (      self  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)     ->				List[Any]:


          '''simple docstring'''







          if self.task in ["default", "seq2seq-lm"]:
              snake_case:						Optional[int]         =   super()._flatten_past_key_values_(__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)
          else:
              snake_case:						List[str]         =   super(__lowerCamelCase  ,		self						)._flatten_past_key_values_(
                  __lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)


