

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 SCREAMING_SNAKE_CASE	:
  """simple docstring"""






  def __init__(			self :    List[str]       ,							lowerCAmelCase :    List[str]       ,							lowerCAmelCase :    Optional[int]=2       ,							lowerCAmelCase :    Tuple=True       ,							lowerCAmelCase :    Union[str, Any]=True       ,							lowerCAmelCase :    List[Any]=10       ,							lowerCAmelCase :    str=10       ,							lowerCAmelCase :    int=8       ,							lowerCAmelCase :    Dict=[2, 4, 4]       ,							lowerCAmelCase :    str=3       ,							lowerCAmelCase :    Dict=7_68       ,							lowerCAmelCase :    List[str]=5       ,							lowerCAmelCase :    str=4       ,							lowerCAmelCase :    List[str]=37       ,							lowerCAmelCase :    Any="gelu_fast"       ,							lowerCAmelCase :    Any=0.0       ,							lowerCAmelCase :    Optional[int]=0.0       ,							lowerCAmelCase :    Optional[int]=0.02       ,							lowerCAmelCase :    Any=1e-06       ,							lowerCAmelCase :    Optional[int]=True       ,							lowerCAmelCase :    Union[str, Any]=None       ,							)							->	Tuple:






     """simple docstring"""







     __lowerCAmelCase  :			Union[str, Any]					    =						parent
     __lowerCAmelCase  :			List[Any]					    =						batch_size
     __lowerCAmelCase  :			Optional[int]					    =						is_training
     __lowerCAmelCase  :			str					    =						use_labels
     __lowerCAmelCase  :			List[str]					    =						num_labels
     __lowerCAmelCase  :			List[str]					    =						image_size
     __lowerCAmelCase  :			List[Any]					    =						num_frames
     __lowerCAmelCase  :			Optional[int]					    =						tubelet_size
     __lowerCAmelCase  :			Union[str, Any]					    =						num_channels
     __lowerCAmelCase  :			List[Any]					    =						hidden_size
     __lowerCAmelCase  :			List[Any]					    =						num_hidden_layers
     __lowerCAmelCase  :			Optional[int]					    =						num_attention_heads
     __lowerCAmelCase  :			str					    =						intermediate_size
     __lowerCAmelCase  :			str					    =						hidden_act
     __lowerCAmelCase  :			Dict					    =						hidden_dropout_prob
     __lowerCAmelCase  :			str					    =						attention_probs_dropout_prob
     __lowerCAmelCase  :			str					    =						initializer_range
     __lowerCAmelCase  :			int					    =						layer_norm_eps
     __lowerCAmelCase  :			str					    =						qkv_bias
     __lowerCAmelCase  :			Optional[int]					    =						scope

     __lowerCAmelCase  :			Any					    =						(
         (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        SCREAMING_SNAKE_CASE    (			self :    Any    )							->	Tuple:






     """simple docstring"""







     __lowerCAmelCase  :			List[Any]					    =						floats_tensor(
         [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]    )

     __lowerCAmelCase  :			int					    =						None
     if self.use_labels:
        __lowerCAmelCase  :			List[str]					    =						ids_tensor([self.batch_size]       ,							self.num_labels    )

     __lowerCAmelCase  :			int					    =						self.get_config()

     return config, pixel_values, labels







  def        SCREAMING_SNAKE_CASE    (			self :    Optional[int]    )							->	Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :			str					    =						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       ,							)
     __lowerCAmelCase  :			List[Any]					    =						self.num_labels
     return config







  def        SCREAMING_SNAKE_CASE    (			self :    Optional[int]       ,							lowerCAmelCase :    Any       ,							lowerCAmelCase :    List[str]       ,							lowerCAmelCase :    str    )							->	int:






     """simple docstring"""







     __lowerCAmelCase  :			Tuple					    =						VivitModel(config=__snake_case    )
     model.to(__snake_case    )
     model.eval()
     __lowerCAmelCase  :			Optional[int]					    =						model(__snake_case    )
     self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)    )







  def        SCREAMING_SNAKE_CASE    (			self :    Dict       ,							lowerCAmelCase :    Tuple       ,							lowerCAmelCase :    int       ,							lowerCAmelCase :    int    )							->	Any:






     """simple docstring"""







     __lowerCAmelCase  :			Any					    =						VivitForVideoClassification(__snake_case    )
     model.to(__snake_case    )
     model.eval()

     __lowerCAmelCase  :			Tuple					    =						model(__snake_case    )

     # verify the logits shape
     __lowerCAmelCase  :			List[Any]					    =						torch.Size((self.batch_size, self.num_labels)    )
     self.parent.assertEqual(result.logits.shape       ,							__snake_case    )







  def        SCREAMING_SNAKE_CASE    (			self :    Union[str, Any]    )							->	List[str]:






     """simple docstring"""







     __lowerCAmelCase  :			List[Any]					    =						self.prepare_config_and_inputs()
     __lowerCAmelCase    ,__lowerCAmelCase    ,__lowerCAmelCase  :			Optional[int]					    =						config_and_inputs
     __lowerCAmelCase  :			List[Any]					    =						{"""pixel_values""": pixel_values}
     return config, inputs_dict






@require_torch
class SCREAMING_SNAKE_CASE	(  _A   ,							unittest.TestCase ):
  """simple docstring"""






  lowerCamelCase    :			List[Any]						      =(VivitModel, VivitForVideoClassification) if is_torch_available() else ()

  lowerCamelCase    :			List[Any]						      =False
  lowerCamelCase    :			List[Any]						      =False
  lowerCamelCase    :			List[Any]						      =False
  lowerCamelCase    :			Optional[Any]						      =False







  def        SCREAMING_SNAKE_CASE    (			self :    List[Any]    )							->	Tuple:






     """simple docstring"""







     __lowerCAmelCase  :			Any					    =						VivitModelTester(self    )
     __lowerCAmelCase  :			List[Any]					    =						ConfigTester(self       ,							config_class=__snake_case       ,							has_text_modality=__snake_case       ,							hidden_size=37    )







  def        SCREAMING_SNAKE_CASE    (			self :    Any       ,							lowerCAmelCase :    Optional[int]       ,							lowerCAmelCase :    Optional[int]       ,							lowerCAmelCase :    int=False    )							->	Tuple:






     """simple docstring"""







     __lowerCAmelCase  :			List[Any]					    =						copy.deepcopy(__snake_case    )

     if return_labels:
        if model_class in get_values(__snake_case    ):
           __lowerCAmelCase  :			Optional[int]					    =						torch.zeros(
               self.model_tester.batch_size       ,							dtype=torch.long       ,							device=__snake_case    )

     return inputs_dict







  def        SCREAMING_SNAKE_CASE    (			self :    int    )							->	List[Any]:






     """simple docstring"""







     self.config_tester.run_common_tests()







  @unittest.skip(reason="""Vivit does not use inputs_embeds"""    )
  def        SCREAMING_SNAKE_CASE    (			self :    Any    )							->	Dict:






     """simple docstring"""







     pass







  def        SCREAMING_SNAKE_CASE    (			self :    Any    )							->	Any:






     """simple docstring"""







     __lowerCAmelCase    ,__lowerCAmelCase  :			Optional[Any]					    =						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        __lowerCAmelCase  :			int					    =						model_class(__snake_case    )
        self.assertIsInstance(model.get_input_embeddings()       ,							(nn.Module)    )
        __lowerCAmelCase  :			List[Any]					    =						model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(__snake_case       ,							nn.Linear    )    )







  def        SCREAMING_SNAKE_CASE    (			self :    Optional[Any]    )							->	Optional[Any]:






     """simple docstring"""







     __lowerCAmelCase    ,__lowerCAmelCase  :			Tuple					    =						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        __lowerCAmelCase  :			Any					    =						model_class(__snake_case    )
        __lowerCAmelCase  :			List[str]					    =						inspect.signature(model.forward    )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        __lowerCAmelCase  :			int					    =						[*signature.parameters.keys()]

        __lowerCAmelCase  :			Any					    =						["""pixel_values""", """head_mask"""]
        self.assertListEqual(arg_names[:2]       ,							__snake_case    )







  def        SCREAMING_SNAKE_CASE    (			self :    Tuple    )							->	Dict:






     """simple docstring"""







     __lowerCAmelCase  :			List[str]					    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__snake_case    )







  def        SCREAMING_SNAKE_CASE    (			self :    int    )							->	int:






     """simple docstring"""







     __lowerCAmelCase  :			str					    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_video_classification(*__snake_case    )







  @slow
  def        SCREAMING_SNAKE_CASE    (			self :    Optional[int]    )							->	Union[str, Any]:






     """simple docstring"""







     for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase  :			str					    =						VivitModel.from_pretrained(__snake_case    )
        self.assertIsNotNone(__snake_case    )







  def        SCREAMING_SNAKE_CASE    (			self :    int    )							->	List[Any]:






     """simple docstring"""







     __lowerCAmelCase    ,__lowerCAmelCase  :			Optional[Any]					    =						self.model_tester.prepare_config_and_inputs_for_common()
     __lowerCAmelCase  :			str					    =						True

     for model_class in self.all_model_classes:
        __lowerCAmelCase  :			Optional[Any]					    =						self.model_tester.seq_length

        __lowerCAmelCase  :			str					    =						True
        __lowerCAmelCase  :			Union[str, Any]					    =						False
        __lowerCAmelCase  :			List[Any]					    =						True
        __lowerCAmelCase  :			int					    =						model_class(__snake_case    )
        model.to(__snake_case    )
        model.eval()
        with torch.no_grad():
           __lowerCAmelCase  :			Dict					    =						model(**self._prepare_for_class(__snake_case       ,							__snake_case    )    )
        __lowerCAmelCase  :			Tuple					    =						outputs.attentions
        self.assertEqual(len(__snake_case    )       ,							self.model_tester.num_hidden_layers    )

        # check that output_attentions also work using config
        del inputs_dict["output_attentions"]
        __lowerCAmelCase  :			Any					    =						True
        __lowerCAmelCase  :			Optional[Any]					    =						model_class(__snake_case    )
        model.to(__snake_case    )
        model.eval()
        with torch.no_grad():
           __lowerCAmelCase  :			Union[str, Any]					    =						model(**self._prepare_for_class(__snake_case       ,							__snake_case    )    )
        __lowerCAmelCase  :			Dict					    =						outputs.attentions
        self.assertEqual(len(__snake_case    )       ,							self.model_tester.num_hidden_layers    )

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

        # Check attention is always last and order is fine
        __lowerCAmelCase  :			Tuple					    =						True
        __lowerCAmelCase  :			Optional[int]					    =						True
        __lowerCAmelCase  :			Dict					    =						model_class(__snake_case    )
        model.to(__snake_case    )
        model.eval()
        with torch.no_grad():
           __lowerCAmelCase  :			List[str]					    =						model(**self._prepare_for_class(__snake_case       ,							__snake_case    )    )

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

        __lowerCAmelCase  :			int					    =						outputs.attentions

        self.assertEqual(len(__snake_case    )       ,							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        SCREAMING_SNAKE_CASE    (			self :    List[str]    )							->	List[Any]:






     """simple docstring"""







     def check_hidden_states_output(lowerCAmelCase :    Optional[Any]       ,							lowerCAmelCase :    str       ,							lowerCAmelCase :    str    ):
        __lowerCAmelCase  :			List[str]					    =						model_class(__snake_case    )
        model.to(__snake_case    )
        model.eval()

        with torch.no_grad():
           __lowerCAmelCase  :			int					    =						model(**self._prepare_for_class(__snake_case       ,							__snake_case    )    )

        __lowerCAmelCase  :			int					    =						outputs.hidden_states
        __lowerCAmelCase  :			str					    =						self.model_tester.num_hidden_layers + 1
        self.assertEqual(len(__snake_case    )       ,							__snake_case    )

        __lowerCAmelCase  :			Optional[int]					    =						self.model_tester.seq_length

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

     __lowerCAmelCase    ,__lowerCAmelCase  :			Tuple					    =						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        __lowerCAmelCase  :			Tuple					    =						True
        check_hidden_states_output(__snake_case       ,							__snake_case       ,							__snake_case    )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        __lowerCAmelCase  :			int					    =						True

        check_hidden_states_output(__snake_case       ,							__snake_case       ,							__snake_case    )







  @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
  def        SCREAMING_SNAKE_CASE    (			self :    int    )							->	List[str]:






     """simple docstring"""







     pass






def 	snake_case_      ()			->	List[Any]:
   __lowerCAmelCase  :			Optional[int]					    =						hf_hub_download(
       repo_id="""hf-internal-testing/spaghetti-video"""					,   filename="""eating_spaghetti_32_frames.npy"""					,   repo_type="""dataset"""     )
   __lowerCAmelCase  :			Tuple					    =						np.load(__UpperCamelCase     )
   return list(__UpperCamelCase     )



@require_torch
@require_vision
class SCREAMING_SNAKE_CASE	(  unittest.TestCase ):
  """simple docstring"""






  @cached_property
  def        SCREAMING_SNAKE_CASE    (			self :    str    )							->	Union[str, Any]:






     """simple docstring"""







     return VivitImageProcessor() if is_vision_available() else None







  @slow
  def        SCREAMING_SNAKE_CASE    (			self :    Any    )							->	Optional[Any]:






     """simple docstring"""







     __lowerCAmelCase  :			List[Any]					    =						VivitForVideoClassification.from_pretrained("""google/vivit-b-16x2-kinetics400"""    ).to(__snake_case    )

     __lowerCAmelCase  :			Optional[Any]					    =						self.default_image_processor
     __lowerCAmelCase  :			Union[str, Any]					    =						prepare_video()
     __lowerCAmelCase  :			str					    =						image_processor(__snake_case       ,							return_tensors="""pt"""    ).to(__snake_case    )

     # forward pass
     with torch.no_grad():
        __lowerCAmelCase  :			Tuple					    =						model(**__snake_case    )

     # verify the logits
     __lowerCAmelCase  :			Dict					    =						torch.Size((1, 4_00)    )
     self.assertEqual(outputs.logits.shape       ,							__snake_case    )

     # taken from original model
     __lowerCAmelCase  :			List[str]					    =						torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]    ).to(__snake_case    )

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








def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						float       ,     __UpperCamelCase			:						float      )  ->							float:
    return a + b


if __name__ == "__main__":
 _lowerCamelCase			 =						5
 _lowerCamelCase			 =						6
 print(F"The sum of {a} + {b} is {add(a, b)}")





