

import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VivitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


if is_torch_available():
       import torch
       from torch import nn

       from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VivitForVideoClassification, VivitModel
       from transformers.models.vivit.modeling_vivit import VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from transformers import VivitImageProcessor



class  __snake_case					:

  def __init__(   self  ,						_A  ,						_A=2  ,						_A=True  ,						_A=True  ,						_A=10  ,						_A=10  ,						_A=8  ,						_A=[2, 4, 4]  ,						_A=3  ,						_A=768  ,						_A=5  ,						_A=4  ,						_A=37  ,						_A="gelu_fast"  ,						_A=0.0  ,						_A=0.0  ,						_A=0.0_2  ,						_A=1E-06  ,						_A=True  ,						_A=None  ,						):
      SCREAMING_SNAKE_CASE_													=  parent
      SCREAMING_SNAKE_CASE_													=  batch_size
      SCREAMING_SNAKE_CASE_													=  is_training
      SCREAMING_SNAKE_CASE_													=  use_labels
      SCREAMING_SNAKE_CASE_													=  num_labels
      SCREAMING_SNAKE_CASE_													=  image_size
      SCREAMING_SNAKE_CASE_													=  num_frames
      SCREAMING_SNAKE_CASE_													=  tubelet_size
      SCREAMING_SNAKE_CASE_													=  num_channels
      SCREAMING_SNAKE_CASE_													=  hidden_size
      SCREAMING_SNAKE_CASE_													=  num_hidden_layers
      SCREAMING_SNAKE_CASE_													=  num_attention_heads
      SCREAMING_SNAKE_CASE_													=  intermediate_size
      SCREAMING_SNAKE_CASE_													=  hidden_act
      SCREAMING_SNAKE_CASE_													=  hidden_dropout_prob
      SCREAMING_SNAKE_CASE_													=  attention_probs_dropout_prob
      SCREAMING_SNAKE_CASE_													=  initializer_range
      SCREAMING_SNAKE_CASE_													=  layer_norm_eps
      SCREAMING_SNAKE_CASE_													=  qkv_bias
      SCREAMING_SNAKE_CASE_													=  scope

      SCREAMING_SNAKE_CASE_													=  (
          (self.image_size // self.tubelet_size[2])
          * (self.image_size // self.tubelet_size[1])
          * (self.num_frames // self.tubelet_size[0])
      ) + 1  # CLS token

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  floats_tensor(
          [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size])

      SCREAMING_SNAKE_CASE_													=  None
      if self.use_labels:
          SCREAMING_SNAKE_CASE_													=  ids_tensor([self.batch_size]  ,						self.num_labels)

      SCREAMING_SNAKE_CASE_													=  self.get_config()

      return config, pixel_values, labels

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  VivitConfig(
          num_frames=self.num_frames  ,						image_size=self.image_size  ,						tubelet_size=self.tubelet_size  ,						num_channels=self.num_channels  ,						hidden_size=self.hidden_size  ,						num_hidden_layers=self.num_hidden_layers  ,						num_attention_heads=self.num_attention_heads  ,						intermediate_size=self.intermediate_size  ,						hidden_act=self.hidden_act  ,						hidden_dropout_prob=self.hidden_dropout_prob  ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,						initializer_range=self.initializer_range  ,						layer_norm_eps=self.layer_norm_eps  ,						qkv_bias=self.qkv_bias  ,						)
      SCREAMING_SNAKE_CASE_													=  self.num_labels
      return config

  def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  VivitModel(config=_A)
      model.to(_A)
      model.eval()
      SCREAMING_SNAKE_CASE_													=  model(_A)
      self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.hidden_size))

  def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A):
      SCREAMING_SNAKE_CASE_													=  VivitForVideoClassification(_A)
      model.to(_A)
      model.eval()

      SCREAMING_SNAKE_CASE_													=  model(_A)

      # verify the logits shape
      SCREAMING_SNAKE_CASE_													=  torch.Size((self.batch_size, self.num_labels))
      self.parent.assertEqual(result.logits.shape  ,						_A)


  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  self.prepare_config_and_inputs()
      SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  config_and_inputs
      SCREAMING_SNAKE_CASE_													=  {'pixel_values': pixel_values}
      return config, inputs_dict






@require_torch
class  __snake_case					(      lowerCAmelCase__       ,	unittest.TestCase    ):
  __lowerCAmelCase       :		Optional[Any]  	=					(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

  __lowerCAmelCase       :		Optional[Any]  	=					False
  __lowerCAmelCase       :		Optional[int]  	=					False
  __lowerCAmelCase       :		int  	=					False
  __lowerCAmelCase       :		Optional[int]  	=					False

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  VivitModelTester(self)
      SCREAMING_SNAKE_CASE_													=  ConfigTester(self  ,						config_class=_A  ,						has_text_modality=_A  ,						hidden_size=37)

  def   lowerCAmelCase__       (   self  ,						_A  ,						_A  ,						_A=False):
      SCREAMING_SNAKE_CASE_													=  copy.deepcopy(_A)

      if return_labels:
          if model_class in get_values(_A):
              SCREAMING_SNAKE_CASE_													=  torch.zeros(
                  self.model_tester.batch_size  ,						dtype=torch.long  ,						device=_A)

      return inputs_dict

  def   lowerCAmelCase__       (   self):
      self.config_tester.run_common_tests()

  @unittest.skip(reason='Vivit does not use inputs_embeds')
  def   lowerCAmelCase__       (   self):
      pass

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
          SCREAMING_SNAKE_CASE_													=  model_class(_A)
          self.assertIsInstance(model.get_input_embeddings()  ,						(nn.Module))
          SCREAMING_SNAKE_CASE_													=  model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(_A  ,						nn.Linear))

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
          SCREAMING_SNAKE_CASE_													=  model_class(_A)
          SCREAMING_SNAKE_CASE_													=  inspect.signature(model.forward)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          SCREAMING_SNAKE_CASE_													=  [*signature.parameters.keys()]

          SCREAMING_SNAKE_CASE_													=  ['pixel_values', 'head_mask']
          self.assertListEqual(arg_names[:2]  ,						_A)

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_A)

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_video_classification(*_A)

  @slow
  def   lowerCAmelCase__       (   self):
      for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          SCREAMING_SNAKE_CASE_													=  VivitModel.from_pretrained(_A)
          self.assertIsNotNone(_A)

  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs_for_common()
      SCREAMING_SNAKE_CASE_													=  True

      for model_class in self.all_model_classes:
          SCREAMING_SNAKE_CASE_													=  self.model_tester.seq_length

          SCREAMING_SNAKE_CASE_													=  True
          SCREAMING_SNAKE_CASE_													=  False
          SCREAMING_SNAKE_CASE_													=  True
          SCREAMING_SNAKE_CASE_													=  model_class(_A)
          model.to(_A)
          model.eval()
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_													=  model(**self._prepare_for_class(_A  ,						_A))
          SCREAMING_SNAKE_CASE_													=  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"]
          SCREAMING_SNAKE_CASE_													=  True
          SCREAMING_SNAKE_CASE_													=  model_class(_A)
          model.to(_A)
          model.eval()
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_													=  model(**self._prepare_for_class(_A  ,						_A))
          SCREAMING_SNAKE_CASE_													=  outputs.attentions
          self.assertEqual(len(_A)  ,						self.model_tester.num_hidden_layers)

          self.assertListEqual(
              list(attentions[0].shape[-3:])  ,						[self.model_tester.num_attention_heads, seq_len, seq_len]  ,						)
          SCREAMING_SNAKE_CASE_													=  len(_A)

          # Check attention is always last and order is fine
          SCREAMING_SNAKE_CASE_													=  True
          SCREAMING_SNAKE_CASE_													=  True
          SCREAMING_SNAKE_CASE_													=  model_class(_A)
          model.to(_A)
          model.eval()
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_													=  model(**self._prepare_for_class(_A  ,						_A))

          self.assertEqual(out_len + 1  ,						len(_A))

          SCREAMING_SNAKE_CASE_													=  outputs.attentions

          self.assertEqual(len(_A)  ,						self.model_tester.num_hidden_layers)

          self.assertListEqual(
              list(self_attentions[0].shape[-3:])  ,						[self.model_tester.num_attention_heads, seq_len, seq_len]  ,						)

  def   lowerCAmelCase__       (   self):
      def check_hidden_states_output(_A  ,						_A  ,						_A):
          SCREAMING_SNAKE_CASE_													=  model_class(_A)
          model.to(_A)
          model.eval()

          with torch.no_grad():
              SCREAMING_SNAKE_CASE_													=  model(**self._prepare_for_class(_A  ,						_A))

          SCREAMING_SNAKE_CASE_													=  outputs.hidden_states
          SCREAMING_SNAKE_CASE_													=  self.model_tester.num_hidden_layers + 1
          self.assertEqual(len(_A)  ,						_A)

          SCREAMING_SNAKE_CASE_													=  self.model_tester.seq_length

          self.assertListEqual(
              list(hidden_states[0].shape[-2:])  ,						[seq_length, self.model_tester.hidden_size]  ,						)

      SCREAMING_SNAKE_CASE_							,     SCREAMING_SNAKE_CASE_													=  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
          SCREAMING_SNAKE_CASE_													=  True
          check_hidden_states_output(_A  ,						_A  ,						_A)

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          SCREAMING_SNAKE_CASE_													=  True

          check_hidden_states_output(_A  ,						_A  ,						_A)


  @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.')
  def   lowerCAmelCase__       (   self):
      pass






def    _UpperCAmelCase    (					):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_													=  hf_hub_download(
        repo_id='hf-internal-testing/spaghetti-video'	,				filename='eating_spaghetti_32_frames.npy'	,				repo_type='dataset'    )
    SCREAMING_SNAKE_CASE_													=  np.load(_SCREAMING_SNAKE_CASE    )
    return list(_SCREAMING_SNAKE_CASE    )





@require_torch
@require_vision
class  __snake_case					(      unittest.TestCase    ):

  @cached_property
  def   lowerCAmelCase__       (   self):
      return VivitImageProcessor() if is_vision_available() else None


  @slow
  def   lowerCAmelCase__       (   self):
      SCREAMING_SNAKE_CASE_													=  VivitForVideoClassification.from_pretrained('google/vivit-b-16x2-kinetics400').to(_A)

      SCREAMING_SNAKE_CASE_													=  self.default_image_processor
      SCREAMING_SNAKE_CASE_													=  prepare_video()
      SCREAMING_SNAKE_CASE_													=  image_processor(_A  ,						return_tensors='pt').to(_A)

      # forward pass
      with torch.no_grad():
          SCREAMING_SNAKE_CASE_													=  model(**_A)

      # verify the logits
      SCREAMING_SNAKE_CASE_													=  torch.Size((1, 400))
      self.assertEqual(outputs.logits.shape  ,						_A)

      # taken from original model
      SCREAMING_SNAKE_CASE_													=  torch.tensor([-1.0_5_4_3, 2.0_7_6_4, -0.2_1_0_4, 0.4_4_3_9, -0.9_6_5_8]).to(_A)

      self.assertTrue(torch.allclose(outputs.logits[0, :5]  ,						_A  ,						atol=1E-4))





from ...configuration_utils import PretrainedConfig
from ...utils import logging


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

UpperCamelCase__					: int								=      {
    "google/pegasus-x-base": "https://huggingface.co/google/pegasus-x-base/resolve/main/config.json",
    "google/pegasus-x-large": "https://huggingface.co/google/pegasus-x-large/resolve/main/config.json",
    # See all PEGASUS-X models at https://huggingface.co/models?filter=pegasus-x
}



class  __snake_case					(      lowerCAmelCase__    ):
  __lowerCAmelCase       :		List[Any]  	=					'pegasus_x'
  __lowerCAmelCase       :		Optional[int]  	=					['past_key_values']
  __lowerCAmelCase       :		List[str]  	=					{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}

  def __init__(   self  ,						_A=96103  ,						_A=16384  ,						_A=16  ,						_A=4096  ,						_A=16  ,						_A=16  ,						_A=4096  ,						_A=16  ,						_A=0.0  ,						_A=0.0  ,						_A=True  ,						_A=True  ,						_A="gelu"  ,						_A=1024  ,						_A=0.1  ,						_A=0.0  ,						_A=0.0  ,						_A=0.0_2  ,						_A=0  ,						_A=True  ,						_A=0  ,						_A=1  ,						_A=1  ,						_A=32  ,						_A=512  ,						_A=True  ,						**_A  ,						):
      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

      SCREAMING_SNAKE_CASE_													=  num_global_tokens
      SCREAMING_SNAKE_CASE_													=  block_size
      SCREAMING_SNAKE_CASE_													=  stagger_local_blocks

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

  @property
  def   lowerCAmelCase__       (   self):
      return self.encoder_attention_heads


  @property
  def   lowerCAmelCase__       (   self):
      return self.d_model


